| """Parsing/inferring signatures from documentation. |
| |
| This module provides several functions to generate better stubs using |
| docstrings and Sphinx docs (.rst files). |
| """ |
| import re |
| import io |
| import contextlib |
| import tokenize |
| |
| from typing import ( |
| Optional, MutableMapping, MutableSequence, List, Sequence, Tuple, NamedTuple, Any |
| ) |
| from typing_extensions import Final |
| |
| # Type alias for signatures strings in format ('func_name', '(arg, opt_arg=False)'). |
| Sig = Tuple[str, str] |
| |
| |
| _TYPE_RE: Final = re.compile(r"^[a-zA-Z_][\w\[\], ]*(\.[a-zA-Z_][\w\[\], ]*)*$") |
| _ARG_NAME_RE: Final = re.compile(r"\**[A-Za-z_][A-Za-z0-9_]*$") |
| |
| |
| def is_valid_type(s: str) -> bool: |
| """Try to determine whether a string might be a valid type annotation.""" |
| if s in ('True', 'False', 'retval'): |
| return False |
| if ',' in s and '[' not in s: |
| return False |
| return _TYPE_RE.match(s) is not None |
| |
| |
| class ArgSig: |
| """Signature info for a single argument.""" |
| |
| def __init__(self, name: str, type: Optional[str] = None, default: bool = False): |
| self.name = name |
| if type and not is_valid_type(type): |
| raise ValueError("Invalid type: " + type) |
| self.type = type |
| # Does this argument have a default value? |
| self.default = default |
| |
| def __repr__(self) -> str: |
| return "ArgSig(name={}, type={}, default={})".format(repr(self.name), repr(self.type), |
| repr(self.default)) |
| |
| def __eq__(self, other: Any) -> bool: |
| if isinstance(other, ArgSig): |
| return (self.name == other.name and self.type == other.type and |
| self.default == other.default) |
| return False |
| |
| |
| class FunctionSig(NamedTuple): |
| name: str |
| args: List[ArgSig] |
| ret_type: str |
| |
| |
| # States of the docstring parser. |
| STATE_INIT: Final = 1 |
| STATE_FUNCTION_NAME: Final = 2 |
| STATE_ARGUMENT_LIST: Final = 3 |
| STATE_ARGUMENT_TYPE: Final = 4 |
| STATE_ARGUMENT_DEFAULT: Final = 5 |
| STATE_RETURN_VALUE: Final = 6 |
| STATE_OPEN_BRACKET: Final = 7 # For generic types. |
| |
| |
| class DocStringParser: |
| """Parse function signatures in documentation.""" |
| |
| def __init__(self, function_name: str) -> None: |
| # Only search for signatures of function with this name. |
| self.function_name = function_name |
| self.state = [STATE_INIT] |
| self.accumulator = "" |
| self.arg_type: Optional[str] = None |
| self.arg_name = "" |
| self.arg_default: Optional[str] = None |
| self.ret_type = "Any" |
| self.found = False |
| self.args: List[ArgSig] = [] |
| # Valid signatures found so far. |
| self.signatures: List[FunctionSig] = [] |
| |
| def add_token(self, token: tokenize.TokenInfo) -> None: |
| """Process next token from the token stream.""" |
| if (token.type == tokenize.NAME and token.string == self.function_name and |
| self.state[-1] == STATE_INIT): |
| self.state.append(STATE_FUNCTION_NAME) |
| |
| elif (token.type == tokenize.OP and token.string == '(' and |
| self.state[-1] == STATE_FUNCTION_NAME): |
| self.state.pop() |
| self.accumulator = "" |
| self.found = True |
| self.state.append(STATE_ARGUMENT_LIST) |
| |
| elif self.state[-1] == STATE_FUNCTION_NAME: |
| # Reset state, function name not followed by '('. |
| self.state.pop() |
| |
| elif (token.type == tokenize.OP and token.string in ('[', '(', '{') and |
| self.state[-1] != STATE_INIT): |
| self.accumulator += token.string |
| self.state.append(STATE_OPEN_BRACKET) |
| |
| elif (token.type == tokenize.OP and token.string in (']', ')', '}') and |
| self.state[-1] == STATE_OPEN_BRACKET): |
| self.accumulator += token.string |
| self.state.pop() |
| |
| elif (token.type == tokenize.OP and token.string == ':' and |
| self.state[-1] == STATE_ARGUMENT_LIST): |
| self.arg_name = self.accumulator |
| self.accumulator = "" |
| self.state.append(STATE_ARGUMENT_TYPE) |
| |
| elif (token.type == tokenize.OP and token.string == '=' and |
| self.state[-1] in (STATE_ARGUMENT_LIST, STATE_ARGUMENT_TYPE)): |
| if self.state[-1] == STATE_ARGUMENT_TYPE: |
| self.arg_type = self.accumulator |
| self.state.pop() |
| else: |
| self.arg_name = self.accumulator |
| self.accumulator = "" |
| self.state.append(STATE_ARGUMENT_DEFAULT) |
| |
| elif (token.type == tokenize.OP and token.string in (',', ')') and |
| self.state[-1] in (STATE_ARGUMENT_LIST, STATE_ARGUMENT_DEFAULT, |
| STATE_ARGUMENT_TYPE)): |
| if self.state[-1] == STATE_ARGUMENT_DEFAULT: |
| self.arg_default = self.accumulator |
| self.state.pop() |
| elif self.state[-1] == STATE_ARGUMENT_TYPE: |
| self.arg_type = self.accumulator |
| self.state.pop() |
| elif self.state[-1] == STATE_ARGUMENT_LIST: |
| self.arg_name = self.accumulator |
| if not (token.string == ')' and self.accumulator.strip() == '') \ |
| and not _ARG_NAME_RE.match(self.arg_name): |
| # Invalid argument name. |
| self.reset() |
| return |
| |
| if token.string == ')': |
| self.state.pop() |
| |
| # 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 |
| self.accumulator = "" |
| |
| elif token.type == tokenize.OP and token.string == '->' and self.state[-1] == STATE_INIT: |
| self.accumulator = "" |
| self.state.append(STATE_RETURN_VALUE) |
| |
| # ENDMAKER is necessary for python 3.4 and 3.5. |
| elif (token.type in (tokenize.NEWLINE, tokenize.ENDMARKER) and |
| self.state[-1] in (STATE_INIT, STATE_RETURN_VALUE)): |
| if self.state[-1] == STATE_RETURN_VALUE: |
| if not is_valid_type(self.accumulator): |
| self.reset() |
| return |
| self.ret_type = self.accumulator |
| self.accumulator = "" |
| self.state.pop() |
| |
| if self.found: |
| self.signatures.append(FunctionSig(name=self.function_name, args=self.args, |
| ret_type=self.ret_type)) |
| self.found = False |
| self.args = [] |
| self.ret_type = 'Any' |
| # Leave state as INIT. |
| else: |
| self.accumulator += token.string |
| |
| def reset(self) -> None: |
| self.state = [STATE_INIT] |
| self.args = [] |
| self.found = False |
| self.accumulator = "" |
| |
| def get_signatures(self) -> List[FunctionSig]: |
| """Return sorted copy of the list of signatures found so far.""" |
| def has_arg(name: str, signature: FunctionSig) -> bool: |
| return any(x.name == name for x in signature.args) |
| |
| def args_kwargs(signature: FunctionSig) -> bool: |
| return has_arg('*args', signature) and has_arg('**kwargs', signature) |
| |
| # Move functions with (*args, **kwargs) in their signature to last place. |
| return list(sorted(self.signatures, key=lambda x: 1 if args_kwargs(x) else 0)) |
| |
| |
| def infer_sig_from_docstring(docstr: Optional[str], name: str) -> Optional[List[FunctionSig]]: |
| """Convert function signature to list of TypedFunctionSig |
| |
| Look for function signatures of function in docstring. Signature is a string of |
| the format <function_name>(<signature>) -> <return type> or perhaps without |
| the return type. |
| |
| Returns empty list, when no signature is found, one signature in typical case, |
| multiple signatures, if docstring specifies multiple signatures for overload functions. |
| Return None if the docstring is empty. |
| |
| Arguments: |
| * docstr: docstring |
| * name: name of function for which signatures are to be found |
| """ |
| if not docstr: |
| return None |
| |
| state = DocStringParser(name) |
| # Return all found signatures, even if there is a parse error after some are found. |
| with contextlib.suppress(tokenize.TokenError): |
| try: |
| tokens = tokenize.tokenize(io.BytesIO(docstr.encode('utf-8')).readline) |
| for token in tokens: |
| state.add_token(token) |
| except IndentationError: |
| return None |
| sigs = state.get_signatures() |
| |
| def is_unique_args(sig: FunctionSig) -> bool: |
| """return true if function argument names are unique""" |
| return len(sig.args) == len({arg.name for arg in sig.args}) |
| |
| # Return only signatures that have unique argument names. Mypy fails on non-unique arg names. |
| return [sig for sig in sigs if is_unique_args(sig)] |
| |
| |
| def infer_arg_sig_from_anon_docstring(docstr: str) -> List[ArgSig]: |
| """Convert signature in form of "(self: TestClass, arg0: str='ada')" to List[TypedArgList].""" |
| ret = infer_sig_from_docstring("stub" + docstr, "stub") |
| if ret: |
| return ret[0].args |
| return [] |
| |
| |
| def infer_ret_type_sig_from_docstring(docstr: str, name: str) -> Optional[str]: |
| """Convert signature in form of "func(self: TestClass, arg0) -> int" to their return type.""" |
| ret = infer_sig_from_docstring(docstr, name) |
| if ret: |
| return ret[0].ret_type |
| return None |
| |
| |
| def infer_ret_type_sig_from_anon_docstring(docstr: str) -> Optional[str]: |
| """Convert signature in form of "(self: TestClass, arg0) -> int" to their return type.""" |
| return infer_ret_type_sig_from_docstring("stub" + docstr.strip(), "stub") |
| |
| |
| def parse_signature(sig: str) -> Optional[Tuple[str, |
| List[str], |
| List[str]]]: |
| """Split function signature into its name, positional an optional arguments. |
| |
| The expected format is "func_name(arg, opt_arg=False)". Return the name of function |
| and lists of positional and optional argument names. |
| """ |
| m = re.match(r'([.a-zA-Z0-9_]+)\(([^)]*)\)', sig) |
| if not m: |
| return None |
| name = m.group(1) |
| name = name.split('.')[-1] |
| arg_string = m.group(2) |
| if not arg_string.strip(): |
| # Simple case -- no arguments. |
| return name, [], [] |
| |
| args = [arg.strip() for arg in arg_string.split(',')] |
| positional = [] |
| optional = [] |
| i = 0 |
| while i < len(args): |
| # Accept optional arguments as in both formats: x=None and [x]. |
| if args[i].startswith('[') or '=' in args[i]: |
| break |
| positional.append(args[i].rstrip('[')) |
| i += 1 |
| if args[i - 1].endswith('['): |
| break |
| while i < len(args): |
| arg = args[i] |
| arg = arg.strip('[]') |
| arg = arg.split('=')[0] |
| optional.append(arg) |
| i += 1 |
| return name, positional, optional |
| |
| |
| def build_signature(positional: Sequence[str], |
| optional: Sequence[str]) -> str: |
| """Build function signature from lists of positional and optional argument names.""" |
| args: MutableSequence[str] = [] |
| args.extend(positional) |
| for arg in optional: |
| if arg.startswith('*'): |
| args.append(arg) |
| else: |
| args.append(f'{arg}=...') |
| sig = f"({', '.join(args)})" |
| # Ad-hoc fixes. |
| sig = sig.replace('(self)', '') |
| return sig |
| |
| |
| def parse_all_signatures(lines: Sequence[str]) -> Tuple[List[Sig], |
| List[Sig]]: |
| """Parse all signatures in a given reST document. |
| |
| Return lists of found signatures for functions and classes. |
| """ |
| sigs = [] |
| class_sigs = [] |
| for line in lines: |
| line = line.strip() |
| m = re.match(r'\.\. *(function|method|class) *:: *[a-zA-Z_]', line) |
| if m: |
| sig = line.split('::')[1].strip() |
| parsed = parse_signature(sig) |
| if parsed: |
| name, fixed, optional = parsed |
| if m.group(1) != 'class': |
| sigs.append((name, build_signature(fixed, optional))) |
| else: |
| class_sigs.append((name, build_signature(fixed, optional))) |
| |
| return sorted(sigs), sorted(class_sigs) |
| |
| |
| def find_unique_signatures(sigs: Sequence[Sig]) -> List[Sig]: |
| """Remove names with duplicate found signatures.""" |
| sig_map: MutableMapping[str, List[str]] = {} |
| for name, sig in sigs: |
| sig_map.setdefault(name, []).append(sig) |
| |
| result = [] |
| for name, name_sigs in sig_map.items(): |
| if len(set(name_sigs)) == 1: |
| result.append((name, name_sigs[0])) |
| return sorted(result) |
| |
| |
| def infer_prop_type_from_docstring(docstr: Optional[str]) -> Optional[str]: |
| """Check for Google/Numpy style docstring type annotation for a property. |
| |
| The docstring has the format "<type>: <descriptions>". |
| In the type string, we allow the following characters: |
| * dot: because sometimes classes are annotated using full path |
| * brackets: to allow type hints like List[int] |
| * comma/space: things like Tuple[int, int] |
| """ |
| if not docstr: |
| return None |
| test_str = r'^([a-zA-Z0-9_, \.\[\]]*): ' |
| m = re.match(test_str, docstr) |
| return m.group(1) if m else None |
