Fixed check for type comments

python · Nov 1, 2020 · ae5771e · ae5771e
1 parent 21e7919
commit ae5771e
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 10 deletions.
diff --git a/mypy/fastparse.py b/mypy/fastparse.py
@@ -32,7 +32,7 @@
 from mypy.types import (
     Type, CallableType, AnyType, UnboundType, TupleType, TypeList, EllipsisType, CallableArgument,
     TypeOfAny, Instance, RawExpressionType, ProperType,
-    UnionType)
+    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,8 @@ def visit(self, node: ast3.expr) -> ProperType: ...
     @overload
     def visit(self, node: Optional[AST]) -> Optional[ProperType]: ...
 
-    def visit(self, node: Optional[AST]) -> Optional[ProperType]:
+    def visit(self, node: Optional[AST],
+              is_type_comment: Optional[bool] = False) -> Optional[ProperType]:
         """Modified visit -- keep track of the stack of nodes"""
         if node is None:
             return None
@@ -1327,6 +1329,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,7 +1426,7 @@ 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:
+    def visit_BinOp(self, n: ast3.BinOp, is_type_comment: Optional[bool] = False) -> Type:
         if not isinstance(n.op, ast3.BitOr):
             return self.invalid_type(n)
 
@@ -1431,7 +1435,7 @@ def visit_BinOp(self, n: ast3.BinOp) -> Type:
         return UnionType([left, right],
                          line=self.line,
                          column=self.convert_column(n.col_offset),
-                         uses_pep604_syntax=True)
+                         pep604_syntax=Pep604Syntax(True, is_type_comment))
 
     def visit_NameConstant(self, n: NameConstant) -> Type:
         if isinstance(n.value, bool):

diff --git a/mypy/typeanal.py b/mypy/typeanal.py
@@ -605,7 +605,9 @@ 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
+        if (t.pep604_syntax
+                and t.pep604_syntax.uses_pep_604_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):

diff --git a/mypy/types.py b/mypy/types.py
@@ -1719,18 +1719,23 @@ def serialize(self) -> JsonDict:
         assert False, "Synthetic types don't serialize"
 
 
+class Pep604Syntax(NamedTuple):
+    uses_pep_604_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,
-                 uses_pep604_syntax: bool = False) -> None:
+                 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.uses_pep604_syntax = uses_pep604_syntax
+        self.pep604_syntax = pep604_syntax
 
     def __hash__(self) -> int:
         return hash(frozenset(self.items))

diff --git a/test-data/unit/check-union-or-syntax.test b/test-data/unit/check-union-or-syntax.test
@@ -61,7 +61,7 @@ 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
+# flags: --python-version 3.6
 x = 1  # type: int | str
 
 [case testUnionOrSyntaxFutureImport]