Skip to content

Commit

Permalink
Fixed stubgen parsing generics from C extensions (#8939)
Browse files Browse the repository at this point in the history
pybind11 is capable of producing type signatures that use generics 
(for example 
https://github.com/pybind/pybind11/blob/4e3d9fea74ed50a042d98f68fa35a3133482289b/include/pybind11/stl.h#L140). 

A user may also opt to write a signature in the docstring that uses generics. Currently 
when stubgen parses one of these generics, it attempts to import a part of it. For 
example if a docstring had 

my_func(str, int) -> List[mypackage.module_being_parsed.MyClass], 

the resulting stub file tries to import List[mypackage.module_being_parsed. 

This change fixes this behaviour by breaking the found type down into the 
multiple types around [], characters, adding any imports from those types 
that are needed, and then stripping out the name of the module being parsed.
  • Loading branch information
AWhetter authored Aug 14, 2020
1 parent 77b8574 commit e4b4959
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 1 deletion.
11 changes: 10 additions & 1 deletion mypy/stubgenc.py
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,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('.')]
Expand Down
75 changes: 75 additions & 0 deletions mypy/test/teststubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -794,6 +794,81 @@ def get_attribute(self) -> None:
generate_c_property_stub('attribute', TestClass.attribute, output, readonly=True)
assert_equal(output, ['@property', 'def attribute(self) -> str: ...'])

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:
Expand Down

0 comments on commit e4b4959

Please sign in to comment.