New semantic analyzer: Fix --disallow-untyped-defs with attrs plugin and forward references

mypy/plugin.py

@@ -39,7 +39,7 @@
 
 import types
 
-from abc import abstractmethod
+from abc import abstractmethod, abstractproperty
 from typing import Any, Callable, List, Tuple, Optional, NamedTuple, TypeVar, Dict
 from mypy_extensions import trait
 
@@ -263,6 +263,14 @@ def defer(self) -> None:
         """
         raise NotImplementedError
 
+    @abstractproperty
+    def final_iteration(self) -> bool:
+        """Is this the final iteration of semantic analysis?
+
+        Only available with new semantic analyzer.
+        """
+        raise NotImplementedError
+
 
 # A context for a function hook that infers the return type of a function with
 # a special signature.

mypy/plugins/attrs.py

@@ -210,6 +210,13 @@ def attr_class_maker_callback(ctx: 'mypy.plugin.ClassDefContext',
 
     attributes = _analyze_class(ctx, auto_attribs, kw_only)
 
+    if ctx.api.options.new_semantic_analyzer:
+        # Check if attribute types are ready.
+        for attr in attributes:
+            if info[attr.name].type is None and not ctx.api.final_iteration:
+                ctx.api.defer()
+                return
+
     # Save the attributes so that subclasses can reuse them.
     ctx.cls.info.metadata['attrs'] = {
         'attributes': [attr.serialize() for attr in attributes],

mypy/plugins/common.py

@@ -88,6 +88,14 @@ def add_method(
     """Adds a new method to a class.
     """
     info = ctx.cls.info
+
+    # First remove any previously generated methods with the same name
+    # to avoid clashes and problems in new semantic analyzer.
+    if name in info.names:
+        sym = info.names[name]
+        if sym.plugin_generated and isinstance(sym.node, FuncDef):
+            ctx.cls.defs.body.remove(sym.node)
+
     self_type = self_type or fill_typevars(info)
     function_type = ctx.api.named_type('__builtins__.function')
 

mypy/semanal.py

@@ -3820,6 +3820,11 @@ def defer(self) -> None:
         assert not self.options.new_semantic_analyzer
         raise NotImplementedError('This is only available with --new-semantic-analyzer')
 
+    @property
+    def final_iteration(self) -> bool:
+        assert not self.options.new_semantic_analyzer
+        raise NotImplementedError('This is only available with --new-semantic-analyzer')
+
 
 def replace_implicit_first_type(sig: FunctionLike, new: Type) -> FunctionLike:
     if isinstance(sig, CallableType):

test-data/unit/check-attr.test

@@ -1023,3 +1023,29 @@ class A(object):
 class B(object):
     x = attr.ib(kw_only=True)  # E: kw_only is not supported in Python 2
 [builtins_py2 fixtures/bool.pyi]
+
+[case testAttrsDisallowUntypedWorksForward]
+# flags: --disallow-untyped-defs
+import attr
+from typing import List
+
+@attr.s
+class B:
+    x: C = attr.ib()
+
+class C(List[C]):
+    pass
+
+reveal_type(B)  # E: Revealed type is 'def (x: __main__.C) -> __main__.B'
+[builtins fixtures/list.pyi]
+
+[case testDisallowUntypedWorksForwardBad]
+# flags: --disallow-untyped-defs
+import attr
+
+@attr.s  # E: Function is missing a type annotation for one or more arguments
+class B:
+    x = attr.ib()  # E: Need type annotation for 'x'
+
+reveal_type(B)  # E: Revealed type is 'def (x: Any) -> __main__.B'
+[builtins fixtures/list.pyi]

test-data/unit/check-dataclasses.test

@@ -559,3 +559,30 @@ class Foo:
     bar: Optional[int] = field(default=None)
 [builtins fixtures/list.pyi]
 [out]
+
+[case testDisallowUntypedWorksForward]
+# flags: --disallow-untyped-defs
+from dataclasses import dataclass
+from typing import List
+
+@dataclass
+class B:
+    x: C
+
+class C(List[C]):
+    pass
+
+reveal_type(B)  # E: Revealed type is 'def (x: __main__.C) -> __main__.B'
+[builtins fixtures/list.pyi]
+
+[case testDisallowUntypedWorksForwardBad]
+# flags: --disallow-untyped-defs
+from dataclasses import dataclass
+
+@dataclass
+class B:
+    x: Undefined  # E: Name 'Undefined' is not defined
+    y = undefined()  # E: Name 'undefined' is not defined
+
+reveal_type(B)  # E: Revealed type is 'def (x: Any) -> __main__.B'
+[builtins fixtures/list.pyi]

test-data/unit/fixtures/list.pyi

@@ -5,7 +5,7 @@ from typing import TypeVar, Generic, Iterable, Iterator, Sequence, overload
 T = TypeVar('T')
 
 class object:
-    def __init__(self): pass
+    def __init__(self) -> None: pass
 
 class type: pass
 class ellipsis: pass