Skip to content

Commit

Permalink
[stubgen] Fix a few bugs with stubgen of c modules (#11096)
Browse files Browse the repository at this point in the history
* [stubgenc] Fix bug when inferring signatures with no args

* [stubgenc] Add self arg when it is omitted from docstring signature

* [stubgenc] Use 'cls' as self arg for classmethods without signature info
  • Loading branch information
chadrik authored Sep 14, 2021
1 parent 3ef8040 commit e654572
Show file tree
Hide file tree
Showing 3 changed files with 55 additions and 10 deletions.
17 changes: 10 additions & 7 deletions mypy/stubdoc.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,13 +148,16 @@ def add_token(self, token: tokenize.TokenInfo) -> None:

if token.string == ')':
self.state.pop()
try:
self.args.append(ArgSig(name=self.arg_name, type=self.arg_type,
default=bool(self.arg_default)))
except ValueError:
# wrong type, use Any
self.args.append(ArgSig(name=self.arg_name, type=None,
default=bool(self.arg_default)))

# arg_name is empty when there are no args. e.g. func()
if self.arg_name:
try:
self.args.append(ArgSig(name=self.arg_name, type=self.arg_type,
default=bool(self.arg_default)))
except ValueError:
# wrong type, use Any
self.args.append(ArgSig(name=self.arg_name, type=None,
default=bool(self.arg_default)))
self.arg_name = ""
self.arg_type = None
self.arg_default = None
Expand Down
11 changes: 8 additions & 3 deletions mypy/stubgenc.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,12 +182,17 @@ def generate_c_function_stub(module: ModuleType,
del inferred[-1]
if not inferred:
if class_name and name not in sigs:
inferred = [FunctionSig(name, args=infer_method_sig(name), ret_type=ret_type)]
inferred = [FunctionSig(name, args=infer_method_sig(name, self_var),
ret_type=ret_type)]
else:
inferred = [FunctionSig(name=name,
args=infer_arg_sig_from_anon_docstring(
sigs.get(name, '(*args, **kwargs)')),
ret_type=ret_type)]
elif class_name and self_var:
args = inferred[0].args
if not args or args[0].name != self_var:
args.insert(0, ArgSig(name=self_var))

is_overloaded = len(inferred) > 1 if inferred else False
if is_overloaded:
Expand Down Expand Up @@ -438,7 +443,7 @@ def is_skipped_attribute(attr: str) -> bool:
)


def infer_method_sig(name: str) -> List[ArgSig]:
def infer_method_sig(name: str, self_var: Optional[str] = None) -> List[ArgSig]:
args: Optional[List[ArgSig]] = None
if name.startswith('__') and name.endswith('__'):
name = name[2:-2]
Expand Down Expand Up @@ -487,4 +492,4 @@ def infer_method_sig(name: str) -> List[ArgSig]:
if args is None:
args = [ArgSig(name='*args'),
ArgSig(name='**kwargs')]
return [ArgSig(name='self')] + args
return [ArgSig(name=self_var or 'self')] + args
37 changes: 37 additions & 0 deletions mypy/test/teststubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,8 @@ def test_find_unique_signatures(self) -> None:
def test_infer_sig_from_docstring(self) -> None:
assert_equal(infer_sig_from_docstring('\nfunc(x) - y', 'func'),
[FunctionSig(name='func', args=[ArgSig(name='x')], ret_type='Any')])
assert_equal(infer_sig_from_docstring('\nfunc(x)', 'func'),
[FunctionSig(name='func', args=[ArgSig(name='x')], ret_type='Any')])

assert_equal(infer_sig_from_docstring('\nfunc(x, Y_a=None)', 'func'),
[FunctionSig(name='func',
Expand Down Expand Up @@ -218,6 +220,13 @@ def test_infer_sig_from_docstring(self) -> None:
[FunctionSig(name='func', args=[ArgSig(name='x', type='int', default=True)],
ret_type='Any')])

assert_equal(infer_sig_from_docstring('\nfunc(x=3)', 'func'),
[FunctionSig(name='func', args=[ArgSig(name='x', type=None, default=True)],
ret_type='Any')])

assert_equal(infer_sig_from_docstring('\nfunc() -> int', 'func'),
[FunctionSig(name='func', args=[], ret_type='int')])

assert_equal(infer_sig_from_docstring('\nfunc(x: int=3) -> int', 'func'),
[FunctionSig(name='func', args=[ArgSig(name='x', type='int', default=True)],
ret_type='int')])
Expand Down Expand Up @@ -737,6 +746,34 @@ def test(self, arg0: str) -> None:
assert_equal(output, ['def test(self, arg0: int) -> Any: ...'])
assert_equal(imports, [])

def test_generate_c_type_with_docstring_no_self_arg(self) -> None:
class TestClass:
def test(self, arg0: str) -> None:
"""
test(arg0: 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: int) -> Any: ...'])
assert_equal(imports, [])

def test_generate_c_type_classmethod(self) -> None:
class TestClass:
@classmethod
def test(cls, arg0: str) -> None:
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='cls', class_name='TestClass')
assert_equal(output, ['def test(cls, *args, **kwargs) -> Any: ...'])
assert_equal(imports, [])

def test_generate_c_type_with_docstring_empty_default(self) -> None:
class TestClass:
def test(self, arg0: str = "") -> None:
Expand Down

0 comments on commit e654572

Please sign in to comment.