Skip to content

Commit

Permalink
Support union type syntax in type applications, aliases and casts
Browse files Browse the repository at this point in the history
These are only available in Python 3.10 mode.

Work on #9880.
  • Loading branch information
JukkaL committed Jul 5, 2021
1 parent 6eafc5e commit 309d797
Show file tree
Hide file tree
Showing 5 changed files with 45 additions and 8 deletions.
7 changes: 5 additions & 2 deletions mypy/exprtotype.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,13 @@

from mypy.nodes import (
Expression, NameExpr, MemberExpr, IndexExpr, RefExpr, TupleExpr, IntExpr, FloatExpr, UnaryExpr,
ComplexExpr, ListExpr, StrExpr, BytesExpr, UnicodeExpr, EllipsisExpr, CallExpr,
ComplexExpr, ListExpr, StrExpr, BytesExpr, UnicodeExpr, EllipsisExpr, CallExpr, OpExpr,
get_member_expr_fullname
)
from mypy.fastparse import parse_type_string
from mypy.types import (
Type, UnboundType, TypeList, EllipsisType, AnyType, CallableArgument, TypeOfAny,
RawExpressionType, ProperType
RawExpressionType, ProperType, UnionType
)


Expand Down Expand Up @@ -150,5 +150,8 @@ def expr_to_unanalyzed_type(expr: Expression, _parent: Optional[Expression] = No
return RawExpressionType(None, 'builtins.complex', line=expr.line, column=expr.column)
elif isinstance(expr, EllipsisExpr):
return EllipsisType(expr.line)
elif isinstance(expr, OpExpr) and expr.op == '|':
return UnionType([expr_to_unanalyzed_type(expr.left),
expr_to_unanalyzed_type(expr.right)])
else:
raise TypeTranslationError()
9 changes: 8 additions & 1 deletion mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -2101,7 +2101,7 @@ def should_wait_rhs(self, rv: Expression) -> bool:
return self.should_wait_rhs(rv.callee)
return False

def can_be_type_alias(self, rv: Expression) -> bool:
def can_be_type_alias(self, rv: Expression, allow_none: bool = False) -> bool:
"""Is this a valid r.h.s. for an alias definition?
Note: this function should be only called for expressions where self.should_wait_rhs()
Expand All @@ -2113,6 +2113,13 @@ def can_be_type_alias(self, rv: Expression) -> bool:
return True
if self.is_none_alias(rv):
return True
if allow_none and isinstance(rv, NameExpr) and rv.fullname == 'builtins.None':
return True
if (isinstance(rv, OpExpr)
and rv.op == '|'
and self.can_be_type_alias(rv.left, allow_none=True)
and self.can_be_type_alias(rv.right, allow_none=True)):
return True
return False

def is_type_ref(self, rv: Expression, bare: bool = False) -> bool:
Expand Down
2 changes: 1 addition & 1 deletion test-data/unit/check-classes.test
Original file line number Diff line number Diff line change
Expand Up @@ -3154,7 +3154,7 @@ def foo(arg: Type[Any]):
from typing import Type, Any
def foo(arg: Type[Any]):
reveal_type(arg.__str__) # N: Revealed type is "def () -> builtins.str"
reveal_type(arg.mro()) # N: Revealed type is "builtins.list[builtins.type]"
reveal_type(arg.mro()) # N: Revealed type is "builtins.list[builtins.type[Any]]"
[builtins fixtures/type.pyi]
[out]

Expand Down
30 changes: 28 additions & 2 deletions test-data/unit/check-union-or-syntax.test
Original file line number Diff line number Diff line change
Expand Up @@ -96,10 +96,22 @@ 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]
T = int | str
x: T
reveal_type(x) # N: Revealed type is "Union[builtins.int, builtins.str]"
S = list[int] | str | None
y: S
reveal_type(y) # N: Revealed type is "Union[builtins.list[builtins.int], builtins.str, None]"
U = str | None
z: U
reveal_type(z) # N: Revealed type is "Union[builtins.str, None]"

def f(): pass

X = int | str | f()
b: X # E: Variable "__main__.X" is not valid as a type \
# N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
[builtins fixtures/type.pyi]


[case testUnionOrSyntaxWithTypeAliasNotAllowed]
Expand Down Expand Up @@ -131,3 +143,17 @@ x: int | None # E: X | Y syntax for unions requires Python 3.10
from lib import x
[file lib.pyi]
x: int | None


[case testUnionOrSyntaxInMiscRuntimeContexts]
# flags: --python-version 3.10
from typing import cast

class C(list[int | None]):
pass

def f() -> object: pass

reveal_type(cast(str | None, f())) # N: Revealed type is "Union[builtins.str, None]"
reveal_type(list[str | None]()) # N: Revealed type is "builtins.list[Union[builtins.str, None]]"
[builtins fixtures/type.pyi]
5 changes: 3 additions & 2 deletions test-data/unit/fixtures/type.pyi
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# builtins stub used in type-related test cases.

from typing import Generic, TypeVar, List
from typing import Generic, TypeVar, List, Union

T = TypeVar('T')

Expand All @@ -10,8 +10,9 @@ class object:

class list(Generic[T]): pass

class type:
class type(Generic[T]):
__name__: str
def __or__(self, other: Union[type, None]) -> type: pass
def mro(self) -> List['type']: pass

class tuple(Generic[T]): pass
Expand Down

0 comments on commit 309d797

Please sign in to comment.