| """Parsing/inferring signatures from documentation. |
| |
| This module provides several functions to generate better stubs using |
| docstrings and Sphinx docs (.rst files). |
| """ |
| |
| from __future__ import annotations |
| |
| import contextlib |
| import io |
| import keyword |
| import re |
| import tokenize |
| from typing import Any, Final, MutableMapping, MutableSequence, NamedTuple, Sequence, Tuple |
| from typing_extensions import TypeAlias as _TypeAlias |
| |
| import mypy.util |
| |
| # Type alias for signatures strings in format ('func_name', '(arg, opt_arg=False)'). |
| Sig: _TypeAlias = 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: str | None = None, |
| *, |
| default: bool = False, |
| default_value: str = "...", |
| ) -> None: |
| self.name = name |
| self.type = type |
| # Does this argument have a default value? |
| self.default = default |
| self.default_value = default_value |
| |
| def is_star_arg(self) -> bool: |
| return self.name.startswith("*") and not self.name.startswith("**") |
| |
| def is_star_kwarg(self) -> bool: |
| return self.name.startswith("**") |
| |
| 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 |
| and self.default_value == other.default_value |
| ) |
| return False |
| |
| |
| class FunctionSig(NamedTuple): |
| name: str |
| args: list[ArgSig] |
| ret_type: str | None |
| |
| def is_special_method(self) -> bool: |
| return bool( |
| self.name.startswith("__") |
| and self.name.endswith("__") |
| and self.args |
| and self.args[0].name in ("self", "cls") |
| ) |
| |
| def has_catchall_args(self) -> bool: |
| """Return if this signature has catchall args: (*args, **kwargs)""" |
| if self.args and self.args[0].name in ("self", "cls"): |
| args = self.args[1:] |
| else: |
| args = self.args |
| return ( |
| len(args) == 2 |
| and all(a.type in (None, "object", "Any", "typing.Any") for a in args) |
| and args[0].is_star_arg() |
| and args[1].is_star_kwarg() |
| ) |
| |
| def is_catchall_signature(self) -> bool: |
| """Return if this signature is the catchall identity: (*args, **kwargs) -> Any""" |
| return self.has_catchall_args() and self.ret_type in (None, "Any", "typing.Any") |
| |
| def format_sig( |
| self, |
| indent: str = "", |
| is_async: bool = False, |
| any_val: str | None = None, |
| docstring: str | None = None, |
| ) -> str: |
| args: list[str] = [] |
| for arg in self.args: |
| arg_def = arg.name |
| |
| if arg_def in keyword.kwlist: |
| arg_def = "_" + arg_def |
| |
| if ( |
| arg.type is None |
| and any_val is not None |
| and arg.name not in ("self", "cls") |
| and not arg.name.startswith("*") |
| ): |
| arg_type: str | None = any_val |
| else: |
| arg_type = arg.type |
| if arg_type: |
| arg_def += ": " + arg_type |
| if arg.default: |
| arg_def += f" = {arg.default_value}" |
| |
| elif arg.default: |
| arg_def += f"={arg.default_value}" |
| |
| args.append(arg_def) |
| |
| retfield = "" |
| ret_type = self.ret_type if self.ret_type else any_val |
| if ret_type is not None: |
| retfield = " -> " + ret_type |
| |
| prefix = "async " if is_async else "" |
| sig = "{indent}{prefix}def {name}({args}){ret}:".format( |
| indent=indent, prefix=prefix, name=self.name, args=", ".join(args), ret=retfield |
| ) |
| if docstring: |
| suffix = f"\n{indent} {mypy.util.quote_docstring(docstring)}" |
| else: |
| suffix = " ..." |
| return f"{sig}{suffix}" |
| |
| |
| # 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: str | None = None |
| self.arg_name = "" |
| self.arg_default: str | None = 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: |
| if self.arg_type and not is_valid_type(self.arg_type): |
| # wrong type, use Any |
| self.args.append( |
| ArgSig(name=self.arg_name, type=None, default=bool(self.arg_default)) |
| ) |
| else: |
| self.args.append( |
| ArgSig( |
| name=self.arg_name, type=self.arg_type, 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 sorted(self.signatures, key=lambda x: 1 if args_kwargs(x) else 0) |
| |
| |
| def infer_sig_from_docstring(docstr: str | None, name: str) -> list[FunctionSig] | None: |
| """Convert function signature to list of FunctionSig |
| |
| 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 (isinstance(docstr, str) and 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) -> str | None: |
| """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) -> str | None: |
| """Convert signature in form of "(self: TestClass, arg0) -> int" to their return type.""" |
| lines = ["stub" + line.strip() for line in docstr.splitlines() if line.strip().startswith("(")] |
| return infer_ret_type_sig_from_docstring("".join(lines), "stub") |
| |
| |
| def parse_signature(sig: str) -> tuple[str, list[str], list[str]] | None: |
| """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: str | None) -> str | None: |
| """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 |
