works for classes

python · Jul 22, 2024 · 1ed6c91 · 1ed6c91
1 parent 95b5ed8
commit 1ed6c91
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 37 deletions.
diff --git a/mypy/stubgenc.py b/mypy/stubgenc.py
@@ -856,6 +856,7 @@ def generate_class_stub(self, class_name: str, cls: type, output: list[str]) ->
                 continue
             elif attr == "__type_params__":
                 inline_generic = generate_inline_generic(value)
+                continue
 
             prop_type_name = self.strip_or_import(self.get_type_annotation(value))
             classvar = self.add_name("typing.ClassVar")

diff --git a/mypy/stubutil.py b/mypy/stubutil.py
@@ -8,8 +8,8 @@
 from abc import abstractmethod
 from collections import defaultdict
 from contextlib import contextmanager
-from typing import Final, Iterable, Iterator, Mapping, ParamSpec, TypeVar, TypeVarTuple, cast
-from typing_extensions import overload
+from typing import Final, Iterable, Iterator, Mapping, TypeVar, cast
+from typing_extensions import ParamSpec, TypeVarTuple, overload
 
 from mypy_extensions import mypyc_attr
 
@@ -845,28 +845,31 @@ def should_reexport(self, name: str, full_module: str, name_is_alias: bool) -> b
 
 
 def generate_inline_generic(type_params: tuple[TypeVar | ParamSpec | TypeVarTuple, ...]) -> str:
-        """Generate stub for inline generic from __type_params__"""
-
-        generic_arg_list: list[str] = []
-
-        for type_param in type_params:
-            # Not done with isinstance checks so compiled code doesn't necessarily need to use
-            # typing.TypeVar, just something with similar duck typing.
-            if hasattr(type_param, "__constraints__"):
-                # Is TypeVar
-                typevar = cast(TypeVar, type_param)
-                if typevar.__bound__:
-                    generic_arg_list.append(f'{typevar.__name__}: {typevar.__bound__}')
-                else:
-                    generic_arg_list.append(f'{typevar.__name__}: {typevar.__constraints__}')
-            elif hasattr(type_param, "__bounds__"):
-                # Is ParamSpec
-                param_spec = cast(ParamSpec, type_param)
-                generic_arg_list.append(f'**{param_spec.__name__}')
+    """Generate stub for inline generic from __type_params__"""
+
+    generic_arg_list: list[str] = []
+
+    for type_param in type_params:
+        # Not done with isinstance checks so compiled code doesn't necessarily need to use
+        # typing.TypeVar, just something with similar duck typing.
+        if hasattr(type_param, "__constraints__"):
+            # Is TypeVar
+            typevar = cast(TypeVar, type_param)
+            if typevar.__bound__:
+                generic_arg_list.append(f"{typevar.__name__}: {typevar.__bound__.__name__}")
+            elif typevar.__constraints__ != ():
+                generic_arg_list.append(
+                    f"{typevar.__name__}: ({', '.join([constraint.__name__ for constraint in typevar.__constraints__])})"
+                )
             else:
-                # Is TypeVarTuple
-                typevar_tuple = cast(TypeVarTuple, type_param)
-                generic_arg_list.append(f'*{typevar_tuple.__name__}')
+                generic_arg_list.append(f"{typevar.__name__}")
+        elif hasattr(type_param, "__bound__"):
+            # Is ParamSpec
+            param_spec = cast(ParamSpec, type_param)
+            generic_arg_list.append(f"**{param_spec.__name__}")
+        else:
+            # Is TypeVarTuple
+            generic_arg_list.append(f"*{type_param.__name__}")
 
-        flat_internals = ",".join(generic_arg_list)
-        return f"[{flat_internals}]"
+    flat_internals = ", ".join(generic_arg_list)
+    return f"[{flat_internals}]"
diff --git a/mypy/test/teststubgen.py b/mypy/test/teststubgen.py
@@ -8,7 +8,8 @@
 import tempfile
 import unittest
 from types import ModuleType
-from typing import Any, ParamSpec, TypeVar, TypeVarTuple
+from typing import Any, TypeVar, cast
+from typing_extensions import ParamSpec, TypeVarTuple
 
 import pytest
 
@@ -826,17 +827,26 @@ def test_infer_cast_sig(self) -> None:
             assert_equal(infer_method_ret_type(f"__{op}__"), op)
 
     def test_generate_inline_generic(self) -> None:
-        T = TypeVar('T')
-        assert generate_inline_generic((T, )) == "[T]"
-        TBound = TypeVar('TBound', bound=int)
-        assert generate_inline_generic((TBound, )) == "[TBound: int]"
-        TBoundTuple = TypeVar('TBoundTuple', int, str)
-        assert generate_inline_generic((TBoundTuple, )) == "[TBoundTuple: (int, str)]"
-        P = ParamSpec('P')
-        assert generate_inline_generic((P, )) == "[**P]"
-        U = TypeVarTuple('U')
-        assert generate_inline_generic((U, )) == "[*U]"
-        assert generate_inline_generic((T, TBound, TBoundTuple, P, U)) =="[T, TBound: int, TBoundTuple: (int, str), **P, *U]"
+        T = TypeVar("T")
+        assert generate_inline_generic((T,)) == "[T]"
+        TBound = TypeVar("TBound", bound=int)
+        assert generate_inline_generic((TBound,)) == "[TBound: int]"
+        TBoundTuple = TypeVar("TBoundTuple", int, str)
+        assert generate_inline_generic((TBoundTuple,)) == "[TBoundTuple: (int, str)]"
+        P = ParamSpec("P")
+        p_tuple = cast(tuple[ParamSpec], (P,))
+        assert generate_inline_generic(p_tuple) == "[**P]"
+        U = TypeVarTuple("U")
+        u_tuple = cast(tuple[TypeVarTuple], (U,))
+        assert generate_inline_generic(u_tuple) == "[*U]"
+        all_tuple = cast(
+            tuple[TypeVar, TypeVar, TypeVar, ParamSpec, TypeVarTuple],
+            (T, TBound, TBoundTuple, P, U),
+        )
+        assert (
+            generate_inline_generic(all_tuple)
+            == "[T, TBound: int, TBoundTuple: (int, str), **P, *U]"
+        )
 
     def test_generate_class_stub_no_crash_for_object(self) -> None:
         output: list[str] = []
@@ -920,6 +930,18 @@ class TestClass(argparse.Action):
         assert_equal(output, ["class C(argparse.Action): ..."])
         assert_equal(gen.get_imports().splitlines(), ["import argparse"])
 
+    def test_inline_generic_class(self) -> None:
+        T = TypeVar("T")
+
+        class TestClass:
+            __type_params__ = (T,)
+
+        output: list[str] = []
+        mod = ModuleType("module", "")
+        gen = InspectionStubGenerator(mod.__name__, known_modules=[mod.__name__], module=mod)
+        gen.generate_class_stub("C", TestClass, output)
+        assert_equal(output, ["class C[T]: ..."])
+
     def test_generate_c_type_inheritance_builtin_type(self) -> None:
         class TestClass(type):
             pass