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

Co-authored-by: Philippe Prados <[email protected]>
Co-authored-by: Shantanu <[email protected]>

Author: 3 people

docs/source/kinds_of_types.rst

@@ -241,6 +241,96 @@ more specific type:
     since the caller may have to use :py:func:`isinstance` before doing anything
     interesting with the value.
 
+.. _alternative_union_syntax:
+
+Alternative union syntax
+------------------------
+
+`PEP 604 <https://www.python.org/dev/peps/pep-0604/>`_ introduced an alternative way
+for writing union types. Starting with **Python 3.10** it is possible to write
+``Union[int, str]`` as ``int | str``. Any of the following options is possible
+
+.. code-block:: python
+
+    from typing import List
+
+    # Use as Union
+    t1: int | str  # equivalent to Union[int, str]
+
+    # Use as Optional
+    t2: int | None  # equivalent to Optional[int]
+
+    # Use in generics
+    t3: List[int | str]  # equivalent to List[Union[int, str]]
+
+    # Use in type aliases
+    T4 = int | None
+    x: T4
+
+    # Quoted variable annotations
+    t5: "int | str"
+
+    # Quoted function annotations
+    def f(t6: "int | str") -> None: ...
+
+    # Type comments
+    t6 = 42  # type: int | str
+
+It is possible to use most of these even for earlier versions. However there are some
+limitations to be aware of.
+
+.. _alternative_union_syntax_stub_files:
+
+Stub files
+""""""""""
+
+All options are supported, regardless of the Python version the project uses.
+
+.. _alternative_union_syntax_37:
+
+Python 3.7 - 3.9
+""""""""""""""""
+
+It is necessary to add ``from __future__ import annotations`` to delay the evaluation
+of type annotations. Not using it would result in a ``TypeError``.
+This does not apply for **type comments**, **quoted function** and **quoted variable** annotations,
+as those also work for earlier versions, see :ref:`below <alternative_union_syntax_older_version>`.
+
+.. warning::
+
+    Type aliases are **NOT** supported! Those result in a ``TypeError`` regardless
+    if the evaluation of type annotations is delayed.
+
+    Dynamic evaluation of annotations is **NOT** possible (e.g. ``typing.get_type_hints`` and ``eval``).
+    See `note PEP 604 <https://www.python.org/dev/peps/pep-0604/#change-only-pep-484-type-hints-to-accept-the-syntax-type1-type2>`_.
+    Use ``typing.Union`` or **Python 3.10** instead if you need those!
+
+.. code-block:: python
+
+    from __future__ import annotations
+
+    t1: int | None
+
+    # Type aliases
+    T2 = int | None  # TypeError!
+
+.. _alternative_union_syntax_older_version:
+
+Older versions
+""""""""""""""
+
++------------------------------------------+-----------+-----------+-----------+
+| Python Version                           | 3.6       | 3.0 - 3.5 | 2.7       |
++==========================================+===========+===========+===========+
+| Type comments                            | yes       | yes       | yes       |
++------------------------------------------+-----------+-----------+-----------+
+| Quoted function annotations              | yes       | yes       |           |
++------------------------------------------+-----------+-----------+-----------+
+| Quoted variable annotations              | yes       |           |           |
++------------------------------------------+-----------+-----------+-----------+
+| Everything else                          |           |           |           |
++------------------------------------------+-----------+-----------+-----------+
+
 .. _strict_optional:
 
 Optional types and the None type

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_evaluated=False).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_evaluated: bool = True,
                  ) -> 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_evaluated = is_evaluated
 
     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=self.is_evaluated,
+                         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

@@ -608,6 +608,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,133 @@
+-- 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 testUnionOrSyntaxWithQuotedFunctionTypes]
+# flags: --python-version 3.4
+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]'
+
+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 testUnionOrSyntaxWithQuotedVariableTypes]
+# flags: --python-version 3.6
+y: "int | str" = 42
+reveal_type(y)  # N: Revealed type is 'Union[builtins.int, builtins.str]'
+
+
+[case testUnionOrSyntaxWithTypeAliasWorking]
+# flags: --python-version 3.10
+from typing import Union
+T = Union[int, str]
+x: T
+reveal_type(x)  # N: Revealed type is 'Union[builtins.int, builtins.str]'
+
+
+[case testUnionOrSyntaxWithTypeAliasNotAllowed]
+# flags: --python-version 3.9
+from __future__ import annotations
+T = int | str  # E: Unsupported left operand type for | ("Type[int]")
+[builtins fixtures/tuple.pyi]
+
+
+[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