Repository: google/pytype

Issue:
Inheritance from generic dataclass deduces wrong __init__
Hello,

This bug is similar to issue #628. When inheriting from a Generic dataclass, pytype fails to deduct the correct __init__ for the new class. This happens within the same file.

Here is a minimal example:

```python
import dataclasses
from typing import Generic TypeVar

K = TypeVar('K')

@dataclasses.dataclass
class BaseClass(Generic[K]):
  arg1: K

@dataclasses.dataclass
class ChildClass(BaseClass[int]):
  pass

ChildClass(arg1=3)
```

this will return the following error:
```
Invalid keyword argument arg1 to function ChildClass.__init__ [wrong-keyword-args]
         Expected: (self)
  Actually passed: (self, arg1)
```

Note however, that this code doesn't throw an error:

```python
import dataclasses
from typing import Generic TypeVar

K = TypeVar('K')

@dataclasses.dataclass
class BaseClass(Generic[K]):
  arg1: K

@dataclasses.dataclass
class ChildClass(BaseClass):
  pass

child = ChildClass(arg1=3)
child.arg1. # type: Any
```
 However in this case the type information regarding arg1 is lost, pytype no longer identifies it as an int.

This last snippet exhibits the desired behaviour, arg1 is correctly identified as an int.

```python
import dataclasses
from typing import Generic TypeVar

K = TypeVar('K')

@dataclasses.dataclass
class BaseClass(Generic[K]):
  arg1: K

@dataclasses.dataclass
class ChildClass(BaseClass[int]):
  pass

child = ChildClass(arg1=3)  # pytype: disable=wrong-keyword-args
child.arg1  # type: int
```
Fix parameterized attributes for dataclass subclasses
In subclasses of a Generic dataclass, the init method accepts the right
argument count, instead of expecting the argument number from its
parent's init method and raising a wrong-arg-count error.

Also add support for checking parameterized types of dataclass and raising
wrong-arg-types.

Fixes https://github.com/google/pytype/issues/1104


Discussion:

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this [failed invocation](https://github.com/google/pytype/pull/1218/checks?check_run_id=6679664704) of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.
Hi, this is a hacky solution for supporting Generic and TypeVar parameters in dataclass inheritance. I have managed to get it working on similar example of https://github.com/google/pytype/issues/1104 . I will add propper tests next.
thanks, this looks good! we'll merge it internally (with your github id) and reexport it.

Resolving patch:
diff --git a/pytype/abstract/class_mixin.py b/pytype/abstract/class_mixin.py
--- a/pytype/abstract/class_mixin.py
+++ b/pytype/abstract/class_mixin.py
@@ -543,6 +543,21 @@ def _get_mro_attrs_for_attrs(self, cls_attrs, metadata_key):
           base_attrs.append(a)
     return base_attrs + cls_attrs
 
+  def _recompute_attrs_type_from_mro(self, all_attrs, type_params):
+    """Traverse the MRO and apply Generic type params to class attributes.
+
+    This IS REQUIRED for dataclass instances that inherits from a Generic.
+
+    Args:
+      all_attrs: All __init__ attributes of a class.
+      type_params: List of ParameterizedClass instances that will override
+        TypeVar attributes in all_attrs.
+    """
+    for typ_name, typ_obj in type_params.items():
+      for attr in all_attrs.values():
+        if typ_name == attr.typ.cls.name:
+          attr.typ = typ_obj
+
   def _get_attrs_from_mro(self, cls_attrs, metadata_key):
     """Traverse the MRO and collect base class attributes for metadata_key."""
 
@@ -552,6 +567,7 @@ def _get_attrs_from_mro(self, cls_attrs, metadata_key):
 
     all_attrs = {}
     sub_attrs = []
+    type_params = {}
     attributes_to_ignore = set()
     for base_cls in reversed(self.mro[1:]):
       if not isinstance(base_cls, Class):
@@ -559,6 +575,12 @@ def _get_attrs_from_mro(self, cls_attrs, metadata_key):
       # Some third-party dataclass implementations add implicit fields that
       # should not be considered inherited attributes.
       attributes_to_ignore.update(getattr(base_cls, "IMPLICIT_FIELDS", ()))
+      # Any subclass of a Parameterized dataclass must inherit attributes from
+      # its parent's init.
+      # See https://github.com/google/pytype/issues/1104
+      if _isinstance(base_cls, "ParameterizedClass"):
+        type_params = base_cls.formal_type_parameters
+        base_cls = base_cls.base_cls
       if metadata_key in base_cls.metadata:
         sub_attrs.append([a for a in base_cls.metadata[metadata_key]
                           if a.name not in attributes_to_ignore])
@@ -566,6 +588,8 @@ def _get_attrs_from_mro(self, cls_attrs, metadata_key):
     for attrs in sub_attrs:
       for a in attrs:
         all_attrs[a.name] = a
+
+    self._recompute_attrs_type_from_mro(all_attrs, type_params)
     return list(all_attrs.values())
 
   def record_attr_ordering(self, own_attrs):


Test patch:
diff --git a/pytype/tests/test_dataclasses.py b/pytype/tests/test_dataclasses.py
--- a/pytype/tests/test_dataclasses.py
+++ b/pytype/tests/test_dataclasses.py
@@ -1057,5 +1057,89 @@ class C(foo.A):
           def __init__(self, x: foo.B, w: int) -> None: ...
       """)
 
+  def test_parameterized_generic(self):
+    ty = self.Infer("""
+      from typing import Generic, TypeVar
+      import dataclasses
+      T = TypeVar('T')
+      @dataclasses.dataclass
+      class Foo(Generic[T]):
+        x: str
+        y: T
+    """)
+    self.assertTypesMatchPytd(ty, """
+      from typing import Dict, Generic, TypeVar, Union
+      import dataclasses
+      T = TypeVar('T')
+      @dataclasses.dataclass
+      class Foo(Generic[T]):
+        x: str
+        y: T
+        __dataclass_fields__: Dict[str, dataclasses.Field[Union[str, T]]]
+        def __init__(self, x: str, y: T) -> None:
+            self = Foo[T]
+    """)
+
+  def test_parameterized_subclass(self):
+    ty = self.Infer("""
+      from typing import Generic, TypeVar
+      import dataclasses
+      T = TypeVar('T')
+      @dataclasses.dataclass
+      class Foo(Generic[T]):
+        x: str
+        y: T
+      @dataclasses.dataclass
+      class Bar(Foo[int]):
+        z: float
+    """)
+    self.assertTypesMatchPytd(ty, """
+      from typing import Dict, Generic, TypeVar, Union
+      import dataclasses
+      T = TypeVar('T')
+      @dataclasses.dataclass
+      class Foo(Generic[T]):
+        x: str
+        y: T
+        __dataclass_fields__: Dict[str, dataclasses.Field[Union[str, T]]]
+        def __init__(self, x: str, y: T) -> None:
+            self = Foo[T]
+      @dataclasses.dataclass
+      class Bar(Foo[int]):
+        z: float
+        __dataclass_fields__: Dict[str, dataclasses.Field[Union[float, int, str]]]
+        def __init__(self, x: str, y: int, z: float) -> None: ...
+    """)
+
+  def test_parameterized_subclass_error_count(self):
+    self.CheckWithErrors("""
+      from typing import Generic, TypeVar
+      import dataclasses
+      T = TypeVar('T')
+      @dataclasses.dataclass
+      class Foo(Generic[T]):
+        x: str
+        y: T
+      @dataclasses.dataclass
+      class Bar(Foo[int]):
+        z: float
+      bar = Bar('test', 10, .4, .4)  # wrong-arg-count
+    """)
+
+  def test_parameterized_subclass_error_type(self):
+    self.CheckWithErrors("""
+      from typing import Generic, TypeVar
+      import dataclasses
+      T = TypeVar('T')
+      @dataclasses.dataclass
+      class Foo(Generic[T]):
+        x: str
+        y: T
+      @dataclasses.dataclass
+      class Bar(Foo[int]):
+        z: float
+      bar = Bar('test', .4, .4)  # wrong-arg-types
+    """)
+
 if __name__ == "__main__":
   test_base.main()