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, Pep604Syntax)
 from mypy import defaults
 from mypy import message_registry, errorcodes as codes
 from mypy.errors import Errors
@@ -241,7 +241,8 @@ def parse_type_comment(type_comment: str,
         converted = TypeConverter(errors,
                                   line=line,
                                   override_column=column,
-                                  assume_str_is_unicode=assume_str_is_unicode).visit(typ.body)
+                                  assume_str_is_unicode=assume_str_is_unicode
+                                  ).visit(typ.body, is_type_comment=True)
         return ignored, converted
 
 
@@ -1318,7 +1319,13 @@ def visit(self, node: ast3.expr) -> ProperType: ...
     @overload
     def visit(self, node: Optional[AST]) -> Optional[ProperType]: ...
 
-    def visit(self, node: Optional[AST]) -> Optional[ProperType]:
+    @overload
+    def visit(self, node: ast3.expr, is_type_comment: bool) -> ProperType: ...
+
+    @overload
+    def visit(self, node: Optional[AST], is_type_comment: bool) -> Optional[ProperType]: ...
+
+    def visit(self, node: Optional[AST], is_type_comment: bool = False) -> Optional[ProperType]:
         """Modified visit -- keep track of the stack of nodes"""
         if node is None:
             return None
@@ -1327,6 +1334,8 @@ def visit(self, node: Optional[AST]) -> Optional[ProperType]:
             method = 'visit_' + node.__class__.__name__
             visitor = getattr(self, method, None)
             if visitor is not None:
+                if visitor == self.visit_BinOp:
+                    return visitor(node, is_type_comment)
                 return visitor(node)
             else:
                 return self.invalid_type(node)
@@ -1422,6 +1431,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, is_type_comment: bool = False) -> 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),
+                         pep604_syntax=Pep604Syntax(True, is_type_comment))
+
     def visit_NameConstant(self, n: NameConstant) -> Type:
         if isinstance(n.value, bool):
             return RawExpressionType(n.value, 'builtins.bool', line=self.line)

mypy/semanal.py

@@ -206,7 +206,7 @@ class SemanticAnalyzer(NodeVisitor[None],
     patches = None  # type: List[Tuple[int, Callable[[], None]]]
     loop_depth = 0         # Depth of breakable loops
     cur_mod_id = ''        # Current module id (or None) (phase 2)
-    is_stub_file = False   # Are we analyzing a stub file?
+    _is_stub_file = False   # Are we analyzing a stub file?
     _is_typeshed_stub_file = False  # Are we analyzing a typeshed stub file?
     imports = None  # type: Set[str]  # Imported modules (during phase 2 analysis)
     # Note: some imports (and therefore dependencies) might
@@ -280,6 +280,10 @@ def __init__(self,
 
     # mypyc doesn't properly handle implementing an abstractproperty
     # with a regular attribute so we make them properties
+    @property
+    def is_stub_file(self) -> bool:
+        return self._is_stub_file
+
     @property
     def is_typeshed_stub_file(self) -> bool:
         return self._is_typeshed_stub_file
@@ -507,7 +511,7 @@ def file_context(self,
         self.cur_mod_node = file_node
         self.cur_mod_id = file_node.fullname
         scope.enter_file(self.cur_mod_id)
-        self.is_stub_file = file_node.path.lower().endswith('.pyi')
+        self._is_stub_file = file_node.path.lower().endswith('.pyi')
         self._is_typeshed_stub_file = is_typeshed_file(file_node.path)
         self.globals = file_node.names
         self.tvar_scope = TypeVarLikeScope()

mypy/semanal_shared.py

@@ -78,6 +78,11 @@ def is_future_flag_set(self, flag: str) -> bool:
         """Is the specific __future__ feature imported"""
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def is_stub_file(self) -> bool:
+        raise NotImplementedError
+
 
 @trait
 class SemanticAnalyzerInterface(SemanticAnalyzerCoreInterface):

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,13 @@ 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.pep604_syntax is not None
+                and t.pep604_syntax.uses_pep604_syntax is True
+                and t.pep604_syntax.is_type_comment is False
+                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("X | Y syntax for unions requires Python 3.10", t)
         return UnionType(self.anal_array(t.items), t.line)
 
     def visit_partial_type(self, t: PartialType) -> Type:

mypy/types.py

@@ -1719,16 +1719,23 @@ def serialize(self) -> JsonDict:
         assert False, "Synthetic types don't serialize"
 
 
+Pep604Syntax = NamedTuple('Pep604Syntax', [
+    ('uses_pep604_syntax', bool),
+    ('is_type_comment', bool)])
+
+
 class UnionType(ProperType):
     """The union type Union[T1, ..., Tn] (at least one type argument)."""
 
-    __slots__ = ('items',)
+    __slots__ = ('items', 'pep604_syntax')
 
-    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,
+                 pep604_syntax: Optional[Pep604Syntax] = None) -> 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.pep604_syntax = pep604_syntax
 
     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.6
+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: X | Y syntax for unions requires Python 3.10
+
+[case testUnionOrSyntaxInStubFile]
+# flags: --python-version 3.6
+from lib import x
+[file lib.pyi]
+x: int | None