Fix [invalid-function-definition] error caused by ignoring dataclasses.KW_ONLY.

rchen152 · rchen152 · commit c6993986b823 · 2024-01-17T17:02:35.000-08:00
* Respects KW_ONLY for dataclass attributes that don't use dataclasses.field.
* Removes unnecessary version check for applying kwonly-ness: it's impossible
  for a user to even set KW_ONLY or kw_only unless they're using a version
  where the field/parameter is available, so once it's been set, we can safely
  assume we're in a high-enough version.

PiperOrigin-RevId: 597970055
diff --git a/pytype/overlays/dataclass_overlay.py b/pytype/overlays/dataclass_overlay.py
@@ -97,7 +97,7 @@ def decorate(self, node, cls):
         continue
       kind = ""
       init = True
-      kw_only = False
+      kw_only = sticky_kwonly
       assert typ
       if match_classvar(typ):
         continue
@@ -112,8 +112,8 @@ def decorate(self, node, cls):
           field = orig.data[0]
           orig = field.default
           init = field.init
-          if self.ctx.python_version >= (3, 10):
-            kw_only = sticky_kwonly if field.kw_only is None else field.kw_only
+          if field.kw_only is not None:
+            kw_only = field.kw_only
 
       if orig and orig.data == [self.ctx.convert.none]:
         # vm._apply_annotation mostly takes care of checking that the default
diff --git a/pytype/tests/test_dataclasses.py b/pytype/tests/test_dataclasses.py
@@ -761,6 +761,26 @@ class A:
         def __init__(self, a1: int, a3: int, *, a2: int = ...) -> None: ...
     """)
 
+  @test_utils.skipBeforePy((3, 10), "KW_ONLY is new in 3.10")
+  def test_kwonly_and_nonfield_default(self):
+    ty = self.Infer("""
+      import dataclasses
+      @dataclasses.dataclass
+      class C:
+        _: dataclasses.KW_ONLY
+        x: int = 0
+        y: str
+    """)
+    self.assertTypesMatchPytd(ty, """
+      import dataclasses
+      @dataclasses.dataclass
+      class C:
+        x: int = ...
+        y: str
+        _: dataclasses.KW_ONLY
+        def __init__(self, *, x: int = ..., y: str) -> None: ...
+    """)
+
   def test_star_import(self):
     with self.DepTree([("foo.pyi", """
       import dataclasses
diff --git a/pytype/tests/test_flax_overlay.py b/pytype/tests/test_flax_overlay.py
@@ -266,6 +266,34 @@ def __init__(
           def replace(self: _TBaz, **kwargs) -> _TBaz: ...
       """)
 
+  @test_utils.skipBeforePy((3, 10), "KW_ONLY is new in 3.10")
+  def test_kwonly(self):
+    with test_utils.Tempdir() as d:
+      self._setup_linen_pyi(d)
+      ty = self.Infer("""
+        import dataclasses
+        from flax import linen as nn
+        class C(nn.Module):
+          _: dataclasses.KW_ONLY
+          x: int = 0
+          y: str
+      """, pythonpath=[d.path])
+    self.assertTypesMatchPytd(ty, """
+      import dataclasses
+      from flax import linen as nn
+      from typing import Any, TypeVar
+
+      _TC = TypeVar('_TC', bound=C)
+
+      @dataclasses.dataclass
+      class C(nn.module.Module):
+        x: int = ...
+        y: str
+        _: dataclasses.KW_ONLY
+        def __init__(self, *, x: int = ..., y: str, name: str = ..., parent: Any = ...) -> None: ...
+        def replace(self: _TC, **kwargs) -> _TC: ...
+    """)
+
 
 if __name__ == "__main__":
   test_base.main()
