Fix dataclass plugin to work with new semantic analyzer (#6515)

Fixes #6324

This enables two test files that previously had few crashes.

Changes in this PR:
* Add `defer()` to public plugin API, otherwise dataclasses will not works with forward references.
* Update one plugin test to use the new analyzer because the output is a bit more verbose, and probably not worth the hassle of making it identical to old analyzer.
* One test is skipped because of existing issue #6454
* Add the generated type variable used for self types to the class body (this is identical to named tuples and `attrs` plugin).

Author: ilevkivskyi

mypy/plugin.py

@@ -253,6 +253,15 @@ def qualified_name(self, n: str) -> str:
         """Make qualified name using current module and enclosing class (if any)."""
         raise NotImplementedError
 
+    @abstractmethod
+    def defer(self) -> None:
+        """Call this to defer the processing of the current node.
+
+        This will request an additional 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/dataclasses.py

@@ -4,7 +4,7 @@
 from mypy.nodes import (
     ARG_OPT, ARG_POS, MDEF, Argument, AssignmentStmt, CallExpr,
     Context, Expression, FuncDef, JsonDict, NameExpr,
-    SymbolTableNode, TempNode, TypeInfo, Var,
+    SymbolTableNode, TempNode, TypeInfo, Var, TypeVarExpr
 )
 from mypy.plugin import ClassDefContext
 from mypy.plugins.common import add_method, _get_decorator_bool_argument
@@ -21,6 +21,8 @@
     'dataclasses.dataclass',
 }  # type: Final
 
+SELF_TVAR_NAME = '_DT'  # type: Final
+
 
 class DataclassAttribute:
     def __init__(
@@ -77,6 +79,12 @@ def transform(self) -> None:
         ctx = self._ctx
         info = self._ctx.cls.info
         attributes = self.collect_attributes()
+        if ctx.api.options.new_semantic_analyzer:
+            # Check if attribute types are ready.
+            for attr in attributes:
+                if info[attr.name].type is None:
+                    ctx.api.defer()
+                    return
         decorator_arguments = {
             'init': _get_decorator_bool_argument(self._ctx, 'init', True),
             'eq': _get_decorator_bool_argument(self._ctx, 'eq', True),
@@ -92,14 +100,23 @@ def transform(self) -> None:
                 return_type=NoneTyp(),
             )
 
+        if (decorator_arguments['eq'] and info.get('__eq__') is None or
+                decorator_arguments['order']):
+            # Type variable for self types in generated methods.
+            obj_type = ctx.api.named_type('__builtins__.object')
+            self_tvar_expr = TypeVarExpr(SELF_TVAR_NAME, info.fullname() + '.' + SELF_TVAR_NAME,
+                                         [], obj_type)
+            info.names[SELF_TVAR_NAME] = SymbolTableNode(MDEF, self_tvar_expr)
+
         # Add an eq method, but only if the class doesn't already have one.
         if decorator_arguments['eq'] and info.get('__eq__') is None:
             for method_name in ['__eq__', '__ne__']:
                 # The TVar is used to enforce that "other" must have
                 # the same type as self (covariant).  Note the
                 # "self_type" parameter to add_method.
                 obj_type = ctx.api.named_type('__builtins__.object')
-                cmp_tvar_def = TypeVarDef('T', 'T', -1, [], obj_type)
+                cmp_tvar_def = TypeVarDef(SELF_TVAR_NAME, info.fullname() + '.' + SELF_TVAR_NAME,
+                                          -1, [], obj_type)
                 cmp_other_type = TypeVarType(cmp_tvar_def)
                 cmp_return_type = ctx.api.named_type('__builtins__.bool')
 
@@ -121,7 +138,8 @@ def transform(self) -> None:
                 # Like for __eq__ and __ne__, we want "other" to match
                 # the self type.
                 obj_type = ctx.api.named_type('__builtins__.object')
-                order_tvar_def = TypeVarDef('T', 'T', -1, [], obj_type)
+                order_tvar_def = TypeVarDef(SELF_TVAR_NAME, info.fullname() + '.' + SELF_TVAR_NAME,
+                                            -1, [], obj_type)
                 order_other_type = TypeVarType(order_tvar_def)
                 order_return_type = ctx.api.named_type('__builtins__.bool')
                 order_args = [

mypy/semanal.py

@@ -3818,6 +3818,10 @@ def add_symbol_table_node(self, name: str, stnode: SymbolTableNode) -> None:
         else:
             self.globals[name] = stnode
 
+    def defer(self) -> None:
+        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):

mypy/test/hacks.py

@@ -9,8 +9,6 @@
 new_semanal_blacklist = [
     'check-async-await.test',
     'check-classes.test',
-    'check-custom-plugin.test',
-    'check-dataclasses.test',
     'check-expressions.test',
     'check-flags.test',
     'check-functions.test',

test-data/unit/check-custom-plugin.test

@@ -491,16 +491,18 @@ class Instr(Generic[T]): ...
 plugins=<ROOT>/test-data/unit/plugins/dyn_class.py
 
 [case testDynamicClassPluginNegatives]
-# flags: --config-file tmp/mypy.ini
+# flags: --new-semantic-analyzer --config-file tmp/mypy.ini
 from mod import declarative_base, Column, Instr, non_declarative_base
 
 Bad1 = non_declarative_base()
 Bad2 = Bad3 = declarative_base()
 
-class C1(Bad1): ...  # E: Invalid base class
-class C2(Bad2): ...  # E: Invalid base class
-class C3(Bad3): ...  # E: Invalid base class
-
+class C1(Bad1): ...  # E: Invalid base class \
+                     # E: Invalid type "__main__.Bad1"
+class C2(Bad2): ...  # E: Invalid base class \
+                     # E: Invalid type "__main__.Bad2"
+class C3(Bad3): ...  # E: Invalid base class \
+                     # E: Invalid type "__main__.Bad3"
 [file mod.py]
 from typing import Generic, TypeVar
 def declarative_base(): ...

test-data/unit/check-dataclasses.test

@@ -393,8 +393,9 @@ class Application:
 
 [builtins fixtures/list.pyi]
 
+-- Blocked by #6454
 [case testDataclassOrderingWithCustomMethods]
-# flags: --python-version 3.6
+# flags: --python-version 3.6  --no-new-semantic-analyzer
 from dataclasses import dataclass
 
 @dataclass(order=True)

test-data/unit/check-incremental.test

@@ -3864,10 +3864,10 @@ class A:
 tmp/b.py:3: error: Revealed type is 'def (a: builtins.int) -> a.A'
 tmp/b.py:4: error: Revealed type is 'def (builtins.object, builtins.object) -> builtins.bool'
 tmp/b.py:5: error: Revealed type is 'def (builtins.object, builtins.object) -> builtins.bool'
-tmp/b.py:6: error: Revealed type is 'def [T] (self: T`-1, other: T`-1) -> builtins.bool'
-tmp/b.py:7: error: Revealed type is 'def [T] (self: T`-1, other: T`-1) -> builtins.bool'
-tmp/b.py:8: error: Revealed type is 'def [T] (self: T`-1, other: T`-1) -> builtins.bool'
-tmp/b.py:9: error: Revealed type is 'def [T] (self: T`-1, other: T`-1) -> builtins.bool'
+tmp/b.py:6: error: Revealed type is 'def [_DT] (self: _DT`-1, other: _DT`-1) -> builtins.bool'
+tmp/b.py:7: error: Revealed type is 'def [_DT] (self: _DT`-1, other: _DT`-1) -> builtins.bool'
+tmp/b.py:8: error: Revealed type is 'def [_DT] (self: _DT`-1, other: _DT`-1) -> builtins.bool'
+tmp/b.py:9: error: Revealed type is 'def [_DT] (self: _DT`-1, other: _DT`-1) -> builtins.bool'
 tmp/b.py:18: error: Unsupported operand types for < ("A" and "int")
 tmp/b.py:19: error: Unsupported operand types for <= ("A" and "int")
 tmp/b.py:20: error: Unsupported operand types for > ("A" and "int")

test-data/unit/deps.test

@@ -1400,6 +1400,7 @@ class B(A):
 
 [out]
 <m.A.(abstract)> -> <m.B.__init__>, m
+<m.A._DT> -> <m.B._DT>
 <m.A.__eq__> -> <m.B.__eq__>
 <m.A.__init__> -> <m.B.__init__>, m.B.__init__
 <m.A.__ne__> -> <m.B.__ne__>