Fixed stubgen parsing generics from C extensions

Author: AWhetter

mypy/stubgenc.py

@@ -201,7 +201,16 @@ def strip_or_import(typ: str, module: ModuleType, imports: List[str]) -> str:
         imports: list of import statements (may be modified during the call)
     """
     stripped_type = typ
-    if module and typ.startswith(module.__name__ + '.'):
+    if any(c in typ for c in '[,'):
+        for subtyp in re.split(r'[\[,\]]', typ):
+            strip_or_import(subtyp.strip(), module, imports)
+        if module:
+            stripped_type = re.sub(
+                r'(^|[\[, ]+)' + re.escape(module.__name__ + '.'),
+                r'\1',
+                typ,
+            )
+    elif module and typ.startswith(module.__name__ + '.'):
         stripped_type = typ[len(module.__name__) + 1:]
     elif '.' in typ:
         arg_module = typ[:typ.rindex('.')]

mypy/test/teststubgen.py

@@ -19,7 +19,9 @@
     mypy_options, is_blacklisted_path, is_non_library_module
 )
 from mypy.stubutil import walk_packages, remove_misplaced_type_comments, common_dir_prefix
-from mypy.stubgenc import generate_c_type_stub, infer_method_sig, generate_c_function_stub
+from mypy.stubgenc import (
+    add_typing_import, generate_c_type_stub, infer_method_sig, generate_c_function_stub,
+)
 from mypy.stubdoc import (
     parse_signature, parse_all_signatures, build_signature, find_unique_signatures,
     infer_sig_from_docstring, infer_prop_type_from_docstring, FunctionSig, ArgSig,
@@ -778,6 +780,81 @@ def test(arg0: str) -> None:
         assert_equal(output, ['def test(arg0: str) -> Action: ...'])
         assert_equal(imports, [])
 
+    def test_generate_c_type_with_single_arg_generic(self) -> None:
+        class TestClass:
+            def test(self, arg0: str) -> None:
+                """
+                test(self: TestClass, arg0: List[int])
+                """
+                pass
+        output = []  # type: List[str]
+        imports = []  # type: List[str]
+        mod = ModuleType(TestClass.__module__, '')
+        generate_c_function_stub(mod, 'test', TestClass.test, output, imports,
+                                 self_var='self', class_name='TestClass')
+        assert_equal(output, ['def test(self, arg0: List[int]) -> Any: ...'])
+        assert_equal(imports, [])
+
+    def test_generate_c_type_with_double_arg_generic(self) -> None:
+        class TestClass:
+            def test(self, arg0: str) -> None:
+                """
+                test(self: TestClass, arg0: Dict[str, int])
+                """
+                pass
+        output = []  # type: List[str]
+        imports = []  # type: List[str]
+        mod = ModuleType(TestClass.__module__, '')
+        generate_c_function_stub(mod, 'test', TestClass.test, output, imports,
+                                 self_var='self', class_name='TestClass')
+        assert_equal(output, ['def test(self, arg0: Dict[str,int]) -> Any: ...'])
+        assert_equal(imports, [])
+
+    def test_generate_c_type_with_nested_generic(self) -> None:
+        class TestClass:
+            def test(self, arg0: str) -> None:
+                """
+                test(self: TestClass, arg0: Dict[str, List[int]])
+                """
+                pass
+        output = []  # type: List[str]
+        imports = []  # type: List[str]
+        mod = ModuleType(TestClass.__module__, '')
+        generate_c_function_stub(mod, 'test', TestClass.test, output, imports,
+                                 self_var='self', class_name='TestClass')
+        assert_equal(output, ['def test(self, arg0: Dict[str,List[int]]) -> Any: ...'])
+        assert_equal(imports, [])
+
+    def test_generate_c_type_with_generic_using_other_module_first(self) -> None:
+        class TestClass:
+            def test(self, arg0: str) -> None:
+                """
+                test(self: TestClass, arg0: Dict[argparse.Action, int])
+                """
+                pass
+        output = []  # type: List[str]
+        imports = []  # type: List[str]
+        mod = ModuleType(TestClass.__module__, '')
+        generate_c_function_stub(mod, 'test', TestClass.test, output, imports,
+                                 self_var='self', class_name='TestClass')
+        assert_equal(output, ['def test(self, arg0: Dict[argparse.Action,int]) -> Any: ...'])
+        assert_equal(imports, ['import argparse'])
+
+    def test_generate_c_type_with_generic_using_other_module_last(self) -> None:
+        class TestClass:
+            def test(self, arg0: str) -> None:
+                """
+                test(self: TestClass, arg0: Dict[str, argparse.Action])
+                """
+                pass
+        output = []  # type: List[str]
+        imports = []  # type: List[str]
+        mod = ModuleType(TestClass.__module__, '')
+        generate_c_function_stub(mod, 'test', TestClass.test, output, imports,
+                                 self_var='self', class_name='TestClass')
+        assert_equal(output, ['def test(self, arg0: Dict[str,argparse.Action]) -> Any: ...'])
+        assert_equal(imports, ['import argparse'])
+
     def test_generate_c_type_with_overload_pybind11(self) -> None:
         class TestClass:
             def __init__(self, arg0: str) -> None: