def infer_method_sig(name: str) -> List[ArgSig]:
if name.startswith('__') and name.endswith('__'):
name = name[2:-2]
if name in ('hash', 'iter', 'next', 'sizeof', 'copy', 'deepcopy',
'reduce', 'getinitargs', 'int', 'float', 'trunc',
'complex', 'bool'):
return []
if name == 'getitem':
return [ArgSig(name='index')]
if name == 'setitem':
return [ArgSig(name='index'), ArgSig(name='object')]
if name in ('delattr', 'getattr'):
return [ArgSig(name='name')]
if name == 'setattr':
return [ArgSig(name='name'), ArgSig(name='value')]
if name == 'getstate':
return []
if name == 'setstate':
return [ArgSig(name='state')]
if name in ('eq', 'ne', 'lt', 'le', 'gt', 'ge', 'add', 'radd', 'sub',
'rsub', 'mul', 'rmul', 'mod', 'rmod', 'floordiv',
'rfloordiv', 'truediv', 'rtruediv', 'divmod', 'rdivmod',
'pow', 'rpow'):
return [ArgSig(name='other')]
if name in ('neg', 'pos'):
return []
return [ArgSig(name='*args'), ArgSig(name='**kwargs')]
def test_repr(self) -> None:
assert_equal(repr(ArgSig(name='asd"dsa')),
"ArgSig(name='asd\"dsa', type=None, default=False)")
assert_equal(repr(ArgSig(name="asd'dsa")),
'ArgSig(name="asd\'dsa", type=None, default=False)')
assert_equal(repr(ArgSig("func", 'str')),
"ArgSig(name='func', type='str', default=False)")
assert_equal(repr(ArgSig("func", 'str', default=True)),
"ArgSig(name='func', type='str', default=True)")
def test_infer_arg_sig_from_docstring(self) -> None:
assert_equal(infer_arg_sig_from_docstring("(*args, **kwargs)"),
[ArgSig(name='*args'), ArgSig(name='**kwargs')])
assert_equal(
infer_arg_sig_from_docstring(
"(x: Tuple[int, Tuple[str, int], str]=(1, ('a', 2), 'y'), y: int=4)"),
[ArgSig(name='x', type='Tuple[int,Tuple[str,int],str]', default=True),
ArgSig(name='y', type='int', default=True)])
def test_infer_sig_from_docstring_duplicate_args(self) -> None:
assert_equal(
infer_sig_from_docstring('\nfunc(x, x) -> str\nfunc(x, y) -> int',
'func'),
[
FunctionSig(name='func',
args=[ArgSig(name='x'),
ArgSig(name='y')],
ret_type='int')
])
def infer_method_sig(name: str) -> List[ArgSig]:
args = None # type: Optional[List[ArgSig]]
if name.startswith('__') and name.endswith('__'):
name = name[2:-2]
if name in ('hash', 'iter', 'next', 'sizeof', 'copy', 'deepcopy',
'reduce', 'getinitargs', 'int', 'float', 'trunc',
'complex', 'bool', 'abs', 'bytes', 'dir', 'len',
'reversed', 'round', 'index', 'enter'):
args = []
elif name == 'getitem':
args = [ArgSig(name='index')]
elif name == 'setitem':
args = [ArgSig(name='index'), ArgSig(name='object')]
elif name in ('delattr', 'getattr'):
args = [ArgSig(name='name')]
elif name == 'setattr':
args = [ArgSig(name='name'), ArgSig(name='value')]
elif name == 'getstate':
args = []
elif name == 'setstate':
args = [ArgSig(name='state')]
elif name in ('eq', 'ne', 'lt', 'le', 'gt', 'ge', 'add', 'radd', 'sub',
'rsub', 'mul', 'rmul', 'mod', 'rmod', 'floordiv',
'rfloordiv', 'truediv', 'rtruediv', 'divmod', 'rdivmod',
'pow', 'rpow', 'xor', 'rxor', 'or', 'ror', 'and', 'rand',
'lshift', 'rlshift', 'rshift', 'rrshift', 'contains',
'delitem', 'iadd', 'iand', 'ifloordiv', 'ilshift',
'imod', 'imul', 'ior', 'ipow', 'irshift', 'isub',
'itruediv', 'ixor'):
args = [ArgSig(name='other')]
elif name in ('neg', 'pos', 'invert'):
args = []
elif name == 'get':
args = [ArgSig(name='instance'), ArgSig(name='owner')]
elif name == 'set':
args = [ArgSig(name='instance'), ArgSig(name='value')]
elif name == 'reduce_ex':
args = [ArgSig(name='protocol')]
elif name == 'exit':
args = [
ArgSig(name='type'),
ArgSig(name='value'),
ArgSig(name='traceback')
]
if args is None:
args = [ArgSig(name='*args'), ArgSig(name='**kwargs')]
return [ArgSig(name='self')] + args
def test_infer_binary_op_sig(self) -> None:
for op in ('eq', 'ne', 'lt', 'le', 'gt', 'ge', 'add', 'radd', 'sub',
'rsub', 'mul', 'rmul'):
assert_equal(infer_method_sig('__%s__' % op),
[self_arg, ArgSig(name='other')])
def test_infer_setitem_sig(self) -> None:
assert_equal(
infer_method_sig('__setitem__'),
[self_arg, ArgSig(name='index'),
ArgSig(name='object')])
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, Y_a=None)', 'func'), [
FunctionSig(
name='func',
args=[ArgSig(name='x'),
ArgSig(name='Y_a', default=True)],
ret_type='Any')
])
assert_equal(infer_sig_from_docstring('\nfunc(x, Y_a=3)', 'func'), [
FunctionSig(
name='func',
args=[ArgSig(name='x'),
ArgSig(name='Y_a', default=True)],
ret_type='Any')
])
assert_equal(
infer_sig_from_docstring('\nfunc(x, Y_a=[1, 2, 3])', 'func'), [
FunctionSig(
name='func',
args=[ArgSig(name='x'),
ArgSig(name='Y_a', default=True)],
ret_type='Any')
])
assert_equal(infer_sig_from_docstring('\nafunc(x) - y', 'func'), [])
assert_equal(infer_sig_from_docstring('\nfunc(x, y', 'func'), [])
assert_equal(infer_sig_from_docstring('\nfunc(x=z(y))', 'func'), [
FunctionSig(name='func',
args=[ArgSig(name='x', default=True)],
ret_type='Any')
])
assert_equal(infer_sig_from_docstring('\nfunc x', 'func'), [])
# Try to infer signature from type annotation.
assert_equal(infer_sig_from_docstring('\nfunc(x: int)', 'func'), [
FunctionSig(name='func',
args=[ArgSig(name='x', type='int')],
ret_type='Any')
])
assert_equal(infer_sig_from_docstring('\nfunc(x: int=3)', 'func'), [
FunctionSig(name='func',
args=[ArgSig(name='x', type='int', default=True)],
ret_type='Any')
])
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')
])
assert_equal(
infer_sig_from_docstring('\nfunc(x: int=3) -> int \n', 'func'), [
FunctionSig(name='func',
args=[ArgSig(name='x', type='int', default=True)],
ret_type='int')
])
assert_equal(
infer_sig_from_docstring('\nfunc(x: Tuple[int, str]) -> str',
'func'),
[
FunctionSig(name='func',
args=[ArgSig(name='x', type='Tuple[int,str]')],
ret_type='str')
])
assert_equal(
infer_sig_from_docstring(
'\nfunc(x: Tuple[int, Tuple[str, int], str], y: int) -> str',
'func'),
[
FunctionSig(name='func',
args=[
ArgSig(name='x',
type='Tuple[int,Tuple[str,int],str]'),
ArgSig(name='y', type='int')
],
ret_type='str')
])
assert_equal(infer_sig_from_docstring('\nfunc(x: foo.bar)', 'func'), [
FunctionSig(name='func',
args=[ArgSig(name='x', type='foo.bar')],
ret_type='Any')
])
assert_equal(
infer_sig_from_docstring('\nfunc(x: list=[1,2,[3,4]])', 'func'), [
FunctionSig(name='func',
args=[ArgSig(name='x', type='list', default=True)],
ret_type='Any')
])
assert_equal(
infer_sig_from_docstring('\nfunc(x: str="nasty[")', 'func'), [
FunctionSig(name='func',
args=[ArgSig(name='x', type='str', default=True)],
ret_type='Any')
])
assert_equal(
infer_sig_from_docstring('\nfunc[(x: foo.bar, invalid]', 'func'),
[])
assert_equal(
infer_sig_from_docstring('\nfunc(x: invalid::type<with_template>)',
'func'),
[
FunctionSig(name='func',
args=[ArgSig(name='x', type=None)],
ret_type='Any')
])
assert_equal(infer_sig_from_docstring('\nfunc(x: str="")', 'func'), [
FunctionSig(name='func',
args=[ArgSig(name='x', type='str', default=True)],
ret_type='Any')
])
if modules:
return modules.group(1).split()
else:
return ['main']
def add_file(self, path: str, result: List[str], header: bool) -> None:
if not os.path.exists(path):
result.append('<%s was not generated>' % path.replace('\\', '/'))
return
if header:
result.append('# {}'.format(path[4:]))
with open(path, encoding='utf8') as file:
result.extend(file.read().splitlines())
self_arg = ArgSig(name='self')
class StubgencSuite(unittest.TestCase):
"""Unit tests for stub generation from C modules using introspection.
Note that these don't cover a lot!
"""
def test_infer_hash_sig(self) -> None:
assert_equal(infer_method_sig('__hash__'), [self_arg])
def test_infer_getitem_sig(self) -> None:
assert_equal(infer_method_sig('__getitem__'),
[self_arg, ArgSig(name='index')])
def test_infer_setitem_sig(self) -> None:
def generate_c_function_stub(
module: ModuleType,
name: str,
obj: object,
output: List[str],
imports: List[str],
self_var: Optional[str] = None,
sigs: Optional[Dict[str, str]] = None,
class_name: Optional[str] = None,
class_sigs: Optional[Dict[str, str]] = None) -> None:
"""Generate stub for a single function or method.
The result (always a single line) will be appended to 'output'.
If necessary, any required names will be added to 'imports'.
The 'class_name' is used to find signature of __init__ or __new__ in
'class_sigs'.
"""
if sigs is None:
sigs = {}
if class_sigs is None:
class_sigs = {}
ret_type = 'None' if name == '__init__' and class_name else 'Any'
if (name in ("__new__", "__init__") and name not in sigs and class_name
and class_name in class_sigs):
inferred: Optional[List[FunctionSig]] = [
FunctionSig(
name=name,
args=infer_arg_sig_from_anon_docstring(class_sigs[class_name]),
ret_type=ret_type,
)
]
else:
docstr = getattr(obj, '__doc__', None)
inferred = infer_sig_from_docstring(docstr, name)
if inferred:
assert docstr is not None
if is_pybind11_overloaded_function_docstring(docstr, name):
# Remove pybind11 umbrella (*args, **kwargs) for overloaded functions
del inferred[-1]
if not inferred:
if class_name and name not in sigs:
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:
imports.append('from typing import overload')
if inferred:
for signature in inferred:
sig = []
for arg in signature.args:
if arg.name == self_var:
arg_def = self_var
else:
arg_def = arg.name
if arg_def == 'None':
arg_def = '_none' # None is not a valid argument name
if arg.type:
arg_def += ": " + strip_or_import(
arg.type, module, imports)
if arg.default:
arg_def += " = ..."
sig.append(arg_def)
if is_overloaded:
output.append('@overload')
output.append('def {function}({args}) -> {ret}: ...'.format(
function=name,
args=", ".join(sig),
ret=strip_or_import(signature.ret_type, module, imports)))
def test_infer_getitem_sig(self) -> None:
assert_equal(infer_method_sig('__getitem__'), [ArgSig(name='index')])
def generate_c_function_stub_costum(
module: ModuleType,
name: str,
obj: object,
output: List[str],
imports: List[str],
self_var: Optional[str] = None,
sigs: Optional[Dict[str, str]] = None,
class_name: Optional[str] = None,
class_sigs: Optional[Dict[str, str]] = None) -> None:
"""Generate stub for a single function or method.
(Custom modified version, prototype is mypy.stubgenc.generate_c_function_stub)
The result (always a single line) will be appended to 'output'.
If necessary, any required names will be added to 'imports'.
The 'class_name' is used to find signature of __init__ or __new__ in
'class_sigs'.
"""
# insert Set type from type for mypy missed it
imports.append("from typing import Set")
if sigs is None:
sigs = {}
if class_sigs is None:
class_sigs = {}
ret_type = 'None' if name == '__init__' and class_name else 'Any'
if (name in ("__new__", "__init__") and name not in sigs and class_name
and class_name in class_sigs):
inferred: Optional[List[FunctionSig]] = [
FunctionSig(
name=name,
args=infer_arg_sig_from_anon_docstring(class_sigs[class_name]),
ret_type=ret_type,
)
]
else:
docstr = getattr(obj, '__doc__', None)
inferred = infer_sig_from_docstring(docstr, name)
if inferred:
assert docstr is not None
if is_pybind11_overloaded_function_docstring(docstr, name):
# Remove pybind11 umbrella (*args, **kwargs) for overloaded functions
del inferred[-1]
if not inferred:
if class_name and name not in sigs:
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:
imports.append('from typing import overload')
#TODO: logic branch too deep, need split
if inferred:
# signature id for overload func, used to pick corresbonding signature from inferred docstring
sigid = 0
for signature in inferred:
arg_sig = []
# in docstring, overload function signature start from 1.
sigid += 1
for arg in signature.args:
if arg.name == self_var:
arg_def = self_var
else:
arg_def = arg.name
if arg_def == 'None':
arg_def = '_none' # None is not a valid argument name
if arg.type:
arg_def += ": " + \
strip_or_import(arg.type, module, imports)
# get function default value from func signature in __doc__
if arg.default:
if is_overloaded:
doc = docstr.split("\n")[3:-1]
for i in range(0, len(doc)):
# get signature from overload function docstr
func_str = refine_func_signature(
doc[i], name, is_overloaded, sigid)
if func_str:
var_str = funcparser.getFuncVarStr(
func_str, arg.name)
default_var = re.search(
r" = .{0,}", var_str)
if default_var:
# parsered default var may contains traill char ",", strip it
arg_def += default_var.group(0).strip(
",")
else:
arg_def += " = ..."
break
else:
# similar like overload function
func_str = refine_func_signature(
docstr.split('\n')[0], name)
var_str = funcparser.getFuncVarStr(
func_str, arg.name)
default_var = re.search(r" = .{0,}", var_str)
if default_var:
arg_def += default_var.group(0).strip(",")
else:
arg_def += " = ..."
arg_sig.append(arg_def)
if is_overloaded:
output.append('@overload')
output.append('def {function}({args}) -> {ret}:'.format(
function=name,
args=", ".join(arg_sig),
ret=strip_or_import(signature.ret_type, module, imports)))
# append function summary from __doc__
output.append(" \"\"\"")
if is_overloaded:
doc = docstr.split("\n")[3:-1]
for i in range(0, len(doc)):
funcsig_reg = re.compile(
str(sigid) + ". " + name + r"\(.*?\) ->.*")
next_funcsig_reg = re.compile(
str(sigid + 1) + ". " + name + r"\(.*?\) ->.*")
if re.match(funcsig_reg, doc[i]):
for j in range(i + 2, len(doc)):
if re.match(next_funcsig_reg, doc[j]):
break
output.append(
' {docline}'.format(docline=doc[j]))
break
else:
funcsig_reg = re.compile(name + r"\(.*?\) ->.*")
for line in docstr.split("\n")[2:-1]:
if re.match(funcsig_reg, line):
continue
output.append(' {docline}'.format(docline=line))
output.append(" \"\"\"")
output.append(" ...\n")
