Add support for union types as X | Y (PEP 604)

mypy/fastparse.py

@@ -31,8 +31,8 @@
 )
 from mypy.types import (
     Type, CallableType, AnyType, UnboundType, TupleType, TypeList, EllipsisType, CallableArgument,
-    TypeOfAny, Instance, RawExpressionType, ProperType
-)
+    TypeOfAny, Instance, RawExpressionType, ProperType,
+    UnionType)
 from mypy import defaults
 from mypy import message_registry, errorcodes as codes
 from mypy.errors import Errors
@@ -1422,6 +1422,17 @@ def _extract_argument_name(self, n: ast3.expr) -> Optional[str]:
     def visit_Name(self, n: Name) -> Type:
         return UnboundType(n.id, line=self.line, column=self.convert_column(n.col_offset))
 
+    def visit_BinOp(self, n: ast3.BinOp) -> Type:
+        if not isinstance(n.op, ast3.BitOr):
+            return self.invalid_type(n)
+
+        left = self.visit(n.left)
+        right = self.visit(n.right)
+        return UnionType([left, right],
+                         line=self.line,
+                         column=self.convert_column(n.col_offset),
+                         is_binary_op=True)
+
     def visit_NameConstant(self, n: NameConstant) -> Type:
         if isinstance(n.value, bool):
             return RawExpressionType(n.value, 'builtins.bool', line=self.line)

mypy/test/testcheck.py

@@ -25,6 +25,7 @@
 # List of files that contain test case descriptions.
 typecheck_files = [
     'check-basic.test',
+    'check-union-or-syntax.test',
     'check-callable.test',
     'check-classes.test',
     'check-statements.test',

mypy/typeanal.py

@@ -605,6 +605,11 @@ def visit_star_type(self, t: StarType) -> Type:
         return StarType(self.anal_type(t.type), t.line)
 
     def visit_union_type(self, t: UnionType) -> Type:
+        if (t.is_binary_op is True
+                and self.api.is_stub_file is False
+                and self.options.python_version < (3, 10)
+                and self.api.is_future_flag_set('annotations') is False):
+            self.fail("Alternative syntax for unions requires Python 3.10 or newer", t)
         return UnionType(self.anal_array(t.items), t.line)
 
     def visit_partial_type(self, t: PartialType) -> Type:

mypy/types.py

@@ -1724,11 +1724,13 @@ class UnionType(ProperType):
 
     __slots__ = ('items',)
 
-    def __init__(self, items: Sequence[Type], line: int = -1, column: int = -1) -> None:
+    def __init__(self, items: Sequence[Type], line: int = -1, column: int = -1,
+                 is_binary_op: bool = False) -> None:
         super().__init__(line, column)
         self.items = flatten_nested_unions(items)
         self.can_be_true = any(item.can_be_true for item in items)
         self.can_be_false = any(item.can_be_false for item in items)
+        self.is_binary_op = is_binary_op
 
     def __hash__(self) -> int:
         return hash(frozenset(self.items))

test-data/unit/check-union-or-syntax.test

@@ -0,0 +1,81 @@
+-- Type checking of union types with '|' syntax
+
+[case testUnionOrSyntaxWithTwoBuiltinsTypes]
+# flags: --python-version 3.10
+from __future__ import annotations
+def f(x: int | str) -> int | str:
+    reveal_type(x)  # N: Revealed type is 'Union[builtins.int, builtins.str]'
+    z: int | str = 0
+    reveal_type(z)  # N: Revealed type is 'Union[builtins.int, builtins.str]'
+    return x
+reveal_type(f)  # N: Revealed type is 'def (x: Union[builtins.int, builtins.str]) -> Union[builtins.int, builtins.str]'
+[builtins fixtures/tuple.pyi]
+
+[case testUnionOrSyntaxWithThreeBuiltinsTypes]
+# flags: --python-version 3.10
+def f(x: int | str | float) -> int | str | float:
+    reveal_type(x)  # N: Revealed type is 'Union[builtins.int, builtins.str, builtins.float]'
+    z: int | str | float = 0
+    reveal_type(z)  # N: Revealed type is 'Union[builtins.int, builtins.str, builtins.float]'
+    return x
+
+reveal_type(f)  # N: Revealed type is 'def (x: Union[builtins.int, builtins.str, builtins.float]) -> Union[builtins.int, builtins.str, builtins.float]'
+
+[case testUnionOrSyntaxWithTwoTypes]
+# flags: --python-version 3.10
+class A: pass
+class B: pass
+def f(x: A | B) -> A | B:
+    reveal_type(x)  # N: Revealed type is 'Union[__main__.A, __main__.B]'
+    z: A | B = A()
+    reveal_type(z)  # N: Revealed type is 'Union[__main__.A, __main__.B]'
+    return x
+reveal_type(f)  # N: Revealed type is 'def (x: Union[__main__.A, __main__.B]) -> Union[__main__.A, __main__.B]'
+
+[case testUnionOrSyntaxWithThreeTypes]
+# flags: --python-version 3.10
+class A: pass
+class B: pass
+class C: pass
+def f(x: A | B | C) -> A | B | C:
+    reveal_type(x)  # N: Revealed type is 'Union[__main__.A, __main__.B, __main__.C]'
+    z: A | B | C = A()
+    reveal_type(z)  # N: Revealed type is 'Union[__main__.A, __main__.B, __main__.C]'
+    return x
+reveal_type(f)  # N: Revealed type is 'def (x: Union[__main__.A, __main__.B, __main__.C]) -> Union[__main__.A, __main__.B, __main__.C]'
+
+[case testUnionOrSyntaxWithLiteral]
+# flags: --python-version 3.10
+from typing_extensions import Literal
+reveal_type(Literal[4] | str)  # N: Revealed type is 'Any'
+[builtins fixtures/tuple.pyi]
+
+[case testUnionOrSyntaxWithBadOperator]
+# flags: --python-version 3.10
+x: 1 + 2  # E: Invalid type comment or annotation
+
+[case testUnionOrSyntaxWithBadOperands]
+# flags: --python-version 3.10
+x: int | 42  # E: Invalid type: try using Literal[42] instead?
+y: 42 | int  # E: Invalid type: try using Literal[42] instead?
+z: str | 42 | int  # E: Invalid type: try using Literal[42] instead?
+
+[case testUnionOrSyntaxInComment]
+# flags: --python-version 3.10
+x = 1  # type: int | str
+
+[case testUnionOrSyntaxFutureImport]
+# flags: --python-version 3.7
+from __future__ import annotations
+x: int | None
+[builtins fixtures/tuple.pyi]
+
+[case testUnionOrSyntaxMissingFutureImport]
+# flags: --python-version 3.9
+x: int | None  # E: Alternative syntax for unions requires Python 3.10 or newer
+
+[case testUnionOrSyntaxInStubFile]
+# flags: --python-version 3.6
+from lib import x
+[file lib.pyi]
+x: int | None