Extend the dataclass plugin to deal with callable properties

aecay · aecay · commit fb0785666243 · 2021-04-07T20:35:53.000+01:00
At runtime, the callable properties of dataclasses are handled in the way one would expect: they are not passed a `self` argument. Mypy, however, just sees them as callable class attributes and generates errors about missing arguments. This is a special case of what is discussed in #708. I donʼt have a general solution for that problem, but for dataclasses, I can fix it by automatically converting the callable entries in a data class into (settable) properties. That makes them work properly via-a-vis the typechecker.
diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py
@@ -13,7 +13,7 @@
     add_method, _get_decorator_bool_argument, deserialize_and_fixup_type,
 )
 from mypy.typeops import map_type_from_supertype
-from mypy.types import Type, Instance, NoneType, TypeVarDef, TypeVarType, get_proper_type
+from mypy.types import Type, Instance, NoneType, TypeVarDef, TypeVarType, CallableType, get_proper_type
 from mypy.server.trigger import make_wildcard_trigger
 
 # The set of decorators that generate dataclasses.
@@ -170,6 +170,8 @@ def transform(self) -> None:
 
         if decorator_arguments['frozen']:
             self._freeze(attributes)
+        else:
+            self._propertize_callables(attributes)
 
         self.reset_init_only_vars(info, attributes)
 
@@ -353,6 +355,24 @@ def _freeze(self, attributes: List[DataclassAttribute]) -> None:
                 var._fullname = info.fullname + '.' + var.name
                 info.names[var.name] = SymbolTableNode(MDEF, var)
 
+    def _propertize_callables(self, attributes: List[DataclassAttribute]) -> None:
+        """Converts all attributes with callable types to @property methods.
+
+        This avoids the typechecker getting confused and thinking that
+        `my_dataclass_instance.callable_attr(foo)` is going to receive a
+        `self` argument (it is not).
+
+        """
+        info = self._ctx.cls.info
+        for attr in attributes:
+            if isinstance(attr.type, CallableType):
+                var = attr.to_var()
+                var.info = info
+                var.is_property = True
+                var.is_settable_property = True
+                var._fullname = info.fullname + '.' + var.name
+                info.names[var.name] = SymbolTableNode(MDEF, var)
+
 
 def dataclass_class_maker_callback(ctx: ClassDefContext) -> None:
     """Hooks into the class typechecking process to add support for dataclasses.
diff --git a/test-data/unit/check-dataclasses.test b/test-data/unit/check-dataclasses.test
@@ -1108,3 +1108,72 @@ class B(A):
 reveal_type(B)  # N: Revealed type is "def (foo: builtins.int) -> __main__.B"
 
 [builtins fixtures/property.pyi]
+
+[case testDataclassCallableProperty]
+# flags: --python-version 3.7
+from dataclasses import dataclass
+from typing import Callable
+
+@dataclass
+class A:
+    foo: Callable[[int], int]
+
+def my_foo(x: int) -> int:
+    return x
+
+a = A(foo=my_foo)
+a.foo(1)
+reveal_type(a.foo)  # N: Revealed type is "def (builtins.int) -> builtins.int"
+reveal_type(A.foo)  # N: Revealed type is "def (builtins.int) -> builtins.int"
+[typing fixtures/typing-medium.pyi]
+[case testDataclassCallableAssignment]
+# flags: --python-version 3.7
+from dataclasses import dataclass
+from typing import Callable
+
+@dataclass
+class A:
+    foo: Callable[[int], int]
+
+def my_foo(x: int) -> int:
+    return x
+
+a = A(foo=my_foo)
+
+def another_foo(x: int) -> int:
+    return x + 1
+
+a.foo = another_foo
+[case testDataclassCallablePropertyWrongType]
+# flags: --python-version 3.7
+from dataclasses import dataclass
+from typing import Callable
+
+@dataclass
+class A:
+    foo: Callable[[int], int]
+
+def my_foo(x: int) -> str:
+    return "foo"
+
+a = A(foo=my_foo)  # E: Argument "foo" to "A" has incompatible type "Callable[[int], str]"; expected "Callable[[int], int]"
+[typing fixtures/typing-medium.pyi]
+[case testDataclassCallablePropertyWrongTypeAssignment]
+# flags: --python-version 3.7
+from dataclasses import dataclass
+from typing import Callable
+
+@dataclass
+class A:
+    foo: Callable[[int], int]
+
+def my_foo(x: int) -> int:
+    return x
+
+a = A(foo=my_foo)
+
+def another_foo(x: int) -> str:
+    return "foo"
+
+a.foo = another_foo  # E: Incompatible types in assignment (expression has type "Callable[[int], str]", variable has type "Callable[[int], int]")
+[typing fixtures/typing-medium.pyi]
