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

mypy/fastparse.py

@@ -31,7 +31,7 @@
 )
 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
@@ -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,
+                                  is_type_comment=True).visit(typ.body)
         return ignored, converted
 
 
@@ -268,6 +269,8 @@ def parse_type_string(expr_string: str, expr_fallback_name: str,
             node.original_str_expr = expr_string
             node.original_str_fallback = expr_fallback_name
             return node
+        elif isinstance(node, UnionType):
+            return node
         else:
             return RawExpressionType(expr_string, expr_fallback_name, line, column)
     except (SyntaxError, ValueError):
@@ -1276,12 +1279,14 @@ def __init__(self,
                  line: int = -1,
                  override_column: int = -1,
                  assume_str_is_unicode: bool = True,
+                 is_type_comment: bool = False,
                  ) -> None:
         self.errors = errors
         self.line = line
         self.override_column = override_column
         self.node_stack = []  # type: List[AST]
         self.assume_str_is_unicode = assume_str_is_unicode
+        self.is_type_comment = is_type_comment
 
     def convert_column(self, column: int) -> int:
         """Apply column override if defined; otherwise return column.
@@ -1422,6 +1427,18 @@ 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_evaluated=(not self.is_type_comment),
+                         uses_pep604_syntax=True)
+
     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,12 @@ 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.uses_pep604_syntax is True
+                and t.is_evaluated 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("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

@@ -1722,13 +1722,18 @@ def serialize(self) -> JsonDict:
 class UnionType(ProperType):
     """The union type Union[T1, ..., Tn] (at least one type argument)."""
 
-    __slots__ = ('items',)
+    __slots__ = ('items', 'is_evaluated', 'uses_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,
+                 is_evaluated: bool = True, uses_pep604_syntax: 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)
+        # is_evaluated should be set to false for type comments and string literals
+        self.is_evaluated = is_evaluated
+        # uses_pep604_syntax is True if Union uses OR syntax (X | Y)
+        self.uses_pep604_syntax = uses_pep604_syntax
 
     def __hash__(self) -> int:
         return hash(frozenset(self.items))

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

@@ -0,0 +1,113 @@
+-- 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 testUnionOrSyntaxWithGenerics]
+# flags: --python-version 3.10
+from typing import List
+x: List[int | str]
+reveal_type(x)  # N: Revealed type is 'builtins.list[Union[builtins.int, builtins.str]]'
+[builtins fixtures/list.pyi]
+
+
+[case testUnionOrSyntaxWithQuotedTypes]
+# flags: --python-version 3.10
+from typing import Union
+def f(x: 'Union[int, str, None]') -> 'Union[int, None]':
+    reveal_type(x)  # N: Revealed type is 'Union[builtins.int, builtins.str, None]'
+    return 42
+reveal_type(f)  # N: Revealed type is 'def (x: Union[builtins.int, builtins.str, None]) -> Union[builtins.int, None]'
+
+# flags: --python-version 3.10
+def g(x: "int | str | None") -> "int | None":
+    reveal_type(x)  # N: Revealed type is 'Union[builtins.int, builtins.str, None]'
+    return 42
+reveal_type(g)  # N: Revealed type is 'def (x: Union[builtins.int, builtins.str, None]) -> Union[builtins.int, None]'
+
+
+[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