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")