Coalesce Literals when printing Unions (#12205)

Instead of printing `Union[Literal[X], Literal[Y]]`, these are now
printed as `Literal[X, Y]`.

Author: intgr

mypy/messages.py

@@ -39,6 +39,7 @@
     IS_SETTABLE, IS_CLASSVAR, IS_CLASS_OR_STATIC,
 )
 from mypy.sametypes import is_same_type
+from mypy.typeops import separate_union_literals
 from mypy.util import unmangle
 from mypy.errorcodes import ErrorCode
 from mypy import message_registry, errorcodes as codes
@@ -1664,6 +1665,16 @@ def format_type_inner(typ: Type,
     def format(typ: Type) -> str:
         return format_type_inner(typ, verbosity, fullnames)
 
+    def format_list(types: Sequence[Type]) -> str:
+        return ', '.join(format(typ) for typ in types)
+
+    def format_literal_value(typ: LiteralType) -> str:
+        if typ.is_enum_literal():
+            underlying_type = format(typ.fallback)
+            return '{}.{}'.format(underlying_type, typ.value)
+        else:
+            return typ.value_repr()
+
     # TODO: show type alias names in errors.
     typ = get_proper_type(typ)
 
@@ -1686,15 +1697,10 @@ def format(typ: Type) -> str:
         elif itype.type.fullname in reverse_builtin_aliases:
             alias = reverse_builtin_aliases[itype.type.fullname]
             alias = alias.split('.')[-1]
-            items = [format(arg) for arg in itype.args]
-            return '{}[{}]'.format(alias, ', '.join(items))
+            return '{}[{}]'.format(alias, format_list(itype.args))
         else:
             # There are type arguments. Convert the arguments to strings.
-            a: List[str] = []
-            for arg in itype.args:
-                a.append(format(arg))
-            s = ', '.join(a)
-            return '{}[{}]'.format(base_str, s)
+            return '{}[{}]'.format(base_str, format_list(itype.args))
     elif isinstance(typ, TypeVarType):
         # This is similar to non-generic instance types.
         return typ.name
@@ -1704,10 +1710,7 @@ def format(typ: Type) -> str:
         # Prefer the name of the fallback class (if not tuple), as it's more informative.
         if typ.partial_fallback.type.fullname != 'builtins.tuple':
             return format(typ.partial_fallback)
-        items = []
-        for t in typ.items:
-            items.append(format(t))
-        s = 'Tuple[{}]'.format(', '.join(items))
+        s = 'Tuple[{}]'.format(format_list(typ.items))
         return s
     elif isinstance(typ, TypedDictType):
         # If the TypedDictType is named, return the name
@@ -1722,24 +1725,34 @@ def format(typ: Type) -> str:
         s = 'TypedDict({{{}}})'.format(', '.join(items))
         return s
     elif isinstance(typ, LiteralType):
-        if typ.is_enum_literal():
-            underlying_type = format(typ.fallback)
-            return 'Literal[{}.{}]'.format(underlying_type, typ.value)
-        else:
-            return str(typ)
+        return 'Literal[{}]'.format(format_literal_value(typ))
     elif isinstance(typ, UnionType):
-        # Only print Unions as Optionals if the Optional wouldn't have to contain another Union
-        print_as_optional = (len(typ.items) -
-                             sum(isinstance(get_proper_type(t), NoneType)
-                                 for t in typ.items) == 1)
-        if print_as_optional:
-            rest = [t for t in typ.items if not isinstance(get_proper_type(t), NoneType)]
-            return 'Optional[{}]'.format(format(rest[0]))
+        literal_items, union_items = separate_union_literals(typ)
+
+        # Coalesce multiple Literal[] members. This also changes output order.
+        # If there's just one Literal item, retain the original ordering.
+        if len(literal_items) > 1:
+            literal_str = 'Literal[{}]'.format(
+                ', '.join(format_literal_value(t) for t in literal_items)
+            )
+
+            if len(union_items) == 1 and isinstance(get_proper_type(union_items[0]), NoneType):
+                return 'Optional[{}]'.format(literal_str)
+            elif union_items:
+                return 'Union[{}, {}]'.format(format_list(union_items), literal_str)
+            else:
+                return literal_str
         else:
-            items = []
-            for t in typ.items:
-                items.append(format(t))
-            s = 'Union[{}]'.format(', '.join(items))
+            # Only print Union as Optional if the Optional wouldn't have to contain another Union
+            print_as_optional = (len(typ.items) -
+                                 sum(isinstance(get_proper_type(t), NoneType)
+                                     for t in typ.items) == 1)
+            if print_as_optional:
+                rest = [t for t in typ.items if not isinstance(get_proper_type(t), NoneType)]
+                return 'Optional[{}]'.format(format(rest[0]))
+            else:
+                s = 'Union[{}]'.format(format_list(typ.items))
+
             return s
     elif isinstance(typ, NoneType):
         return 'None'

mypy/typeops.py

@@ -846,3 +846,18 @@ def is_redundant_literal_instance(general: ProperType, specific: ProperType) ->
         return True
 
     return False
+
+
+def separate_union_literals(t: UnionType) -> Tuple[Sequence[LiteralType], Sequence[Type]]:
+    """Separate literals from other members in a union type."""
+    literal_items = []
+    union_items = []
+
+    for item in t.items:
+        proper = get_proper_type(item)
+        if isinstance(proper, LiteralType):
+            literal_items.append(proper)
+        else:
+            union_items.append(item)
+
+    return literal_items, union_items

test-data/unit/check-enum.test

@@ -1405,9 +1405,9 @@ class E(Enum):
 
 e: E
 a: Literal[E.A, E.B, E.C] = e
-b: Literal[E.A, E.B] = e  # E: Incompatible types in assignment (expression has type "E", variable has type "Union[Literal[E.A], Literal[E.B]]")
-c: Literal[E.A, E.C] = e  # E: Incompatible types in assignment (expression has type "E", variable has type "Union[Literal[E.A], Literal[E.C]]")
-b = a  # E: Incompatible types in assignment (expression has type "Union[Literal[E.A], Literal[E.B], Literal[E.C]]", variable has type "Union[Literal[E.A], Literal[E.B]]")
+b: Literal[E.A, E.B] = e  # E: Incompatible types in assignment (expression has type "E", variable has type "Literal[E.A, E.B]")
+c: Literal[E.A, E.C] = e  # E: Incompatible types in assignment (expression has type "E", variable has type "Literal[E.A, E.C]")
+b = a  # E: Incompatible types in assignment (expression has type "Literal[E.A, E.B, E.C]", variable has type "Literal[E.A, E.B]")
 [builtins fixtures/bool.pyi]
 
 [case testIntEnumWithNewTypeValue]

test-data/unit/check-expressions.test

@@ -2153,9 +2153,9 @@ def returns_1_or_2() -> Literal[1, 2]:
     ...
 THREE: Final = 3
 
-if returns_a_or_b() == 'c':  # E: Non-overlapping equality check (left operand type: "Union[Literal['a'], Literal['b']]", right operand type: "Literal['c']")
+if returns_a_or_b() == 'c':  # E: Non-overlapping equality check (left operand type: "Literal['a', 'b']", right operand type: "Literal['c']")
     ...
-if returns_1_or_2() is THREE:  # E: Non-overlapping identity check (left operand type: "Union[Literal[1], Literal[2]]", right operand type: "Literal[3]")
+if returns_1_or_2() is THREE:  # E: Non-overlapping identity check (left operand type: "Literal[1, 2]", right operand type: "Literal[3]")
     ...
 [builtins fixtures/bool.pyi]
 

test-data/unit/check-isinstance.test

@@ -2390,8 +2390,8 @@ class B:
     def t3(self) -> None:
         if isinstance(self, (A1, A2)):
             reveal_type(self)           # N: Revealed type is "Union[__main__.<subclass of "A1" and "B">2, __main__.<subclass of "A2" and "B">]"
-            x0: Literal[0] = self.f()   # E: Incompatible types in assignment (expression has type "Union[Literal[1], Literal[2]]", variable has type "Literal[0]")
-            x1: Literal[1] = self.f()   # E: Incompatible types in assignment (expression has type "Union[Literal[1], Literal[2]]", variable has type "Literal[1]")
+            x0: Literal[0] = self.f()   # E: Incompatible types in assignment (expression has type "Literal[1, 2]", variable has type "Literal[0]")
+            x1: Literal[1] = self.f()   # E: Incompatible types in assignment (expression has type "Literal[1, 2]", variable has type "Literal[1]")
 
 [builtins fixtures/isinstance.pyi]
 

test-data/unit/check-literal.test

@@ -846,7 +846,7 @@ def func(x: Literal['foo', 'bar', '  foo  ']) -> None: ...
 func('foo')
 func('bar')
 func('  foo  ')
-func('baz')       # E: Argument 1 to "func" has incompatible type "Literal['baz']"; expected "Union[Literal['foo'], Literal['bar'], Literal['  foo  ']]"
+func('baz')       # E: Argument 1 to "func" has incompatible type "Literal['baz']"; expected "Literal['foo', 'bar', '  foo  ']"
 
 a: Literal['foo']
 b: Literal['bar']
@@ -860,7 +860,7 @@ func(b)
 func(c)
 func(d)
 func(e)
-func(f)  # E: Argument 1 to "func" has incompatible type "Union[Literal['foo'], Literal['bar'], Literal['baz']]"; expected "Union[Literal['foo'], Literal['bar'], Literal['  foo  ']]"
+func(f)  # E: Argument 1 to "func" has incompatible type "Literal['foo', 'bar', 'baz']"; expected "Literal['foo', 'bar', '  foo  ']"
 [builtins fixtures/tuple.pyi]
 [out]
 
@@ -1129,8 +1129,8 @@ d: int
 
 foo(a)
 foo(b)
-foo(c)  # E: Argument 1 to "foo" has incompatible type "Union[Literal[4], Literal[5]]"; expected "Union[Literal[1], Literal[2], Literal[3]]"
-foo(d)  # E: Argument 1 to "foo" has incompatible type "int"; expected "Union[Literal[1], Literal[2], Literal[3]]"
+foo(c)  # E: Argument 1 to "foo" has incompatible type "Literal[4, 5]"; expected "Literal[1, 2, 3]"
+foo(d)  # E: Argument 1 to "foo" has incompatible type "int"; expected "Literal[1, 2, 3]"
 [builtins fixtures/tuple.pyi]
 [out]
 
@@ -1144,7 +1144,7 @@ c: Literal[4, 'foo']
 
 foo(a)
 foo(b)
-foo(c)  # E: Argument 1 to "foo" has incompatible type "Union[Literal[4], Literal['foo']]"; expected "int"
+foo(c)  # E: Argument 1 to "foo" has incompatible type "Literal[4, 'foo']"; expected "int"
 [builtins fixtures/tuple.pyi]
 [out]
 
@@ -1248,19 +1248,19 @@ class Contravariant(Generic[T_contra]): pass
 a1: Invariant[Literal[1]]
 a2: Invariant[Literal[1, 2]]
 a3: Invariant[Literal[1, 2, 3]]
-a2 = a1  # E: Incompatible types in assignment (expression has type "Invariant[Literal[1]]", variable has type "Invariant[Union[Literal[1], Literal[2]]]")
-a2 = a3  # E: Incompatible types in assignment (expression has type "Invariant[Union[Literal[1], Literal[2], Literal[3]]]", variable has type "Invariant[Union[Literal[1], Literal[2]]]")
+a2 = a1  # E: Incompatible types in assignment (expression has type "Invariant[Literal[1]]", variable has type "Invariant[Literal[1, 2]]")
+a2 = a3  # E: Incompatible types in assignment (expression has type "Invariant[Literal[1, 2, 3]]", variable has type "Invariant[Literal[1, 2]]")
 
 b1: Covariant[Literal[1]]
 b2: Covariant[Literal[1, 2]]
 b3: Covariant[Literal[1, 2, 3]]
 b2 = b1
-b2 = b3  # E: Incompatible types in assignment (expression has type "Covariant[Union[Literal[1], Literal[2], Literal[3]]]", variable has type "Covariant[Union[Literal[1], Literal[2]]]")
+b2 = b3  # E: Incompatible types in assignment (expression has type "Covariant[Literal[1, 2, 3]]", variable has type "Covariant[Literal[1, 2]]")
 
 c1: Contravariant[Literal[1]]
 c2: Contravariant[Literal[1, 2]]
 c3: Contravariant[Literal[1, 2, 3]]
-c2 = c1  # E: Incompatible types in assignment (expression has type "Contravariant[Literal[1]]", variable has type "Contravariant[Union[Literal[1], Literal[2]]]")
+c2 = c1  # E: Incompatible types in assignment (expression has type "Contravariant[Literal[1]]", variable has type "Contravariant[Literal[1, 2]]")
 c2 = c3
 [builtins fixtures/tuple.pyi]
 [out]
@@ -1275,12 +1275,12 @@ def bar(x: Sequence[Literal[1, 2]]) -> None: pass
 a: List[Literal[1]]
 b: List[Literal[1, 2, 3]]
 
-foo(a)  # E: Argument 1 to "foo" has incompatible type "List[Literal[1]]"; expected "List[Union[Literal[1], Literal[2]]]" \
+foo(a)  # E: Argument 1 to "foo" has incompatible type "List[Literal[1]]"; expected "List[Literal[1, 2]]" \
         # N: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance \
         # N: Consider using "Sequence" instead, which is covariant
-foo(b)  # E: Argument 1 to "foo" has incompatible type "List[Union[Literal[1], Literal[2], Literal[3]]]"; expected "List[Union[Literal[1], Literal[2]]]"
+foo(b)  # E: Argument 1 to "foo" has incompatible type "List[Literal[1, 2, 3]]"; expected "List[Literal[1, 2]]"
 bar(a)
-bar(b)  # E: Argument 1 to "bar" has incompatible type "List[Union[Literal[1], Literal[2], Literal[3]]]"; expected "Sequence[Union[Literal[1], Literal[2]]]"
+bar(b)  # E: Argument 1 to "bar" has incompatible type "List[Literal[1, 2, 3]]"; expected "Sequence[Literal[1, 2]]"
 [builtins fixtures/list.pyi]
 [out]
 
@@ -1363,9 +1363,9 @@ x = b  # E: Incompatible types in assignment (expression has type "str", variabl
 y = c  # E: Incompatible types in assignment (expression has type "bool", variable has type "Literal[True]")
 z = d  # This is ok: Literal[None] and None are equivalent.
 
-combined = a  # E: Incompatible types in assignment (expression has type "int", variable has type "Union[Literal[1], Literal['foo'], Literal[True], None]")
-combined = b  # E: Incompatible types in assignment (expression has type "str", variable has type "Union[Literal[1], Literal['foo'], Literal[True], None]")
-combined = c  # E: Incompatible types in assignment (expression has type "bool", variable has type "Union[Literal[1], Literal['foo'], Literal[True], None]")
+combined = a  # E: Incompatible types in assignment (expression has type "int", variable has type "Optional[Literal[1, 'foo', True]]")
+combined = b  # E: Incompatible types in assignment (expression has type "str", variable has type "Optional[Literal[1, 'foo', True]]")
+combined = c  # E: Incompatible types in assignment (expression has type "bool", variable has type "Optional[Literal[1, 'foo', True]]")
 combined = d  # Also ok, for similar reasons.
 
 e: Literal[1] = 1
@@ -1392,9 +1392,9 @@ a: Literal[1] = 2           # E: Incompatible types in assignment (expression ha
 b: Literal["foo"] = "bar"   # E: Incompatible types in assignment (expression has type "Literal['bar']", variable has type "Literal['foo']")
 c: Literal[True] = False    # E: Incompatible types in assignment (expression has type "Literal[False]", variable has type "Literal[True]")
 
-d: Literal[1, 2] = 3                # E: Incompatible types in assignment (expression has type "Literal[3]", variable has type "Union[Literal[1], Literal[2]]")
-e: Literal["foo", "bar"] = "baz"    # E: Incompatible types in assignment (expression has type "Literal['baz']", variable has type "Union[Literal['foo'], Literal['bar']]")
-f: Literal[True, 4] = False         # E: Incompatible types in assignment (expression has type "Literal[False]", variable has type "Union[Literal[True], Literal[4]]")
+d: Literal[1, 2] = 3                # E: Incompatible types in assignment (expression has type "Literal[3]", variable has type "Literal[1, 2]")
+e: Literal["foo", "bar"] = "baz"    # E: Incompatible types in assignment (expression has type "Literal['baz']", variable has type "Literal['foo', 'bar']")
+f: Literal[True, 4] = False         # E: Incompatible types in assignment (expression has type "Literal[False]", variable has type "Literal[True, 4]")
 
 [builtins fixtures/primitives.pyi]
 [out]
@@ -1504,7 +1504,7 @@ reveal_type(arr3)  # N: Revealed type is "builtins.list[builtins.int*]"
 reveal_type(arr4)  # N: Revealed type is "builtins.list[builtins.object*]"
 reveal_type(arr5)  # N: Revealed type is "builtins.list[builtins.object*]"
 
-bad: List[Literal[1, 2]] = [1, 2, 3]  # E: List item 2 has incompatible type "Literal[3]"; expected "Union[Literal[1], Literal[2]]"
+bad: List[Literal[1, 2]] = [1, 2, 3]  # E: List item 2 has incompatible type "Literal[3]"; expected "Literal[1, 2]"
 
 [builtins fixtures/list.pyi]
 [out]
@@ -1619,19 +1619,19 @@ reveal_type(func(a))  # N: Revealed type is "Union[__main__.A, __main__.C]"
 reveal_type(func(b))  # N: Revealed type is "__main__.B"
 reveal_type(func(c))  # N: Revealed type is "Union[__main__.B, __main__.A]"
 reveal_type(func(d))  # N: Revealed type is "__main__.B" \
-                      # E: Argument 1 to "func" has incompatible type "Union[Literal[6], Literal[7]]"; expected "Union[Literal[3], Literal[4], Literal[5], Literal[6]]"
+                      # E: Argument 1 to "func" has incompatible type "Literal[6, 7]"; expected "Literal[3, 4, 5, 6]"
 
 reveal_type(func(e))  # E: No overload variant of "func" matches argument type "int" \
                       # N: Possible overload variants: \
                       # N:     def func(x: Literal[-40]) -> A \
-                      # N:     def func(x: Union[Literal[3], Literal[4], Literal[5], Literal[6]]) -> B \
+                      # N:     def func(x: Literal[3, 4, 5, 6]) -> B \
                       # N:     def func(x: Literal['foo']) -> C \
                       # N: Revealed type is "Any"
 
-reveal_type(func(f))  # E: No overload variant of "func" matches argument type "Union[Literal[7], Literal['bar']]" \
+reveal_type(func(f))  # E: No overload variant of "func" matches argument type "Literal[7, 'bar']" \
                       # N: Possible overload variants: \
                       # N:     def func(x: Literal[-40]) -> A \
-                      # N:     def func(x: Union[Literal[3], Literal[4], Literal[5], Literal[6]]) -> B \
+                      # N:     def func(x: Literal[3, 4, 5, 6]) -> B \
                       # N:     def func(x: Literal['foo']) -> C \
                       # N: Revealed type is "Any"
 [builtins fixtures/tuple.pyi]
@@ -1657,7 +1657,7 @@ reveal_type(f(1))  # N: Revealed type is "builtins.int"
 reveal_type(f(2))  # N: Revealed type is "builtins.int"
 reveal_type(f(y))  # N: Revealed type is "builtins.object"
 reveal_type(f(z))  # N: Revealed type is "builtins.int" \
-                   # E: Argument 1 to "f" has incompatible type "Union[Literal[1], Literal[2], Literal['three']]"; expected "Union[Literal[1], Literal[2]]"
+                   # E: Argument 1 to "f" has incompatible type "Literal[1, 2, 'three']"; expected "Literal[1, 2]"
 [builtins fixtures/tuple.pyi]
 [out]
 
@@ -1726,8 +1726,8 @@ def f(x):
 
 x: Literal['a', 'b']
 y: Literal['a', 'b']
-f(x, y)  # E: Argument 1 to "f" has incompatible type "Union[Literal['a'], Literal['b']]"; expected "Literal['a']" \
-         # E: Argument 2 to "f" has incompatible type "Union[Literal['a'], Literal['b']]"; expected "Literal['a']" \
+f(x, y)  # E: Argument 1 to "f" has incompatible type "Literal['a', 'b']"; expected "Literal['a']" \
+         # E: Argument 2 to "f" has incompatible type "Literal['a', 'b']"; expected "Literal['a']" \
 [builtins fixtures/tuple.pyi]
 [out]
 

test-data/unit/check-narrowing.test

@@ -894,7 +894,7 @@ else:
     reveal_type(y)   # N: Revealed type is "__main__.Custom"
 
 # No contamination here
-if 1 == x == z:      # E: Non-overlapping equality check (left operand type: "Union[Literal[1], Literal[2], None]", right operand type: "Default")
+if 1 == x == z:      # E: Non-overlapping equality check (left operand type: "Optional[Literal[1, 2]]", right operand type: "Default")
     reveal_type(x)   # E: Statement is unreachable
     reveal_type(z)
 else: