function.py

"""
    sphinx.pycode.function
    ~~~~~~~~~~~~~~~~~~~~~~

    Utilities parsing function definitions

    :copyright: Copyright 2007-2020 by the Sphinx team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

import ast
import sys
from collections import OrderedDict
from inspect import Parameter
from typing import cast
from typing import Mapping

if sys.version_info > (3, 8):
    from functools import cached_property
else:
    from cached_property import cached_property


def stringify(node: ast.AST) -> str:
    """Stringify an AST node."""
    if node is None:
        return None
    elif isinstance(node, ast.Attribute):
        return '%s.%s' % (stringify(node.value), node.attr)
    elif isinstance(node, ast.Bytes):
        return repr(node.s)
    elif isinstance(node, ast.Call):
        args = ([stringify(e) for e in node.args] +
                ['%s=%s' % (k.arg, stringify(k.value)) for k in node.keywords])
        return '%s(%s)' % (stringify(node.func), ', '.join(args))
    elif isinstance(node, ast.Dict):
        keys = (stringify(k) for k in node.keys)
        values = (stringify(v) for v in node.values)
        items = (k + ': ' + v for k, v in zip(keys, values))
        return '{' + ', '.join(items) + '}'
    elif isinstance(node, ast.Ellipsis):
        return '...'
    elif sys.version_info > (3, 6) and isinstance(node, ast.Constant):
        return repr(node.value)
    elif isinstance(node, ast.Index):
        return stringify(node.value)
    elif isinstance(node, ast.Lambda):
        return '<function <lambda>>'  # TODO
    elif isinstance(node, ast.List):
        return '[' + ', '.join(stringify(e) for e in node.elts) + ']'
    elif isinstance(node, ast.Name):
        return node.id
    elif isinstance(node, ast.NameConstant):
        return repr(node.value)
    elif isinstance(node, ast.Num):
        return repr(node.n)
    elif isinstance(node, ast.Set):
        return "{" + ', '.join(stringify(e) for e in node.elts) + "}"
    elif isinstance(node, ast.Str):
        return repr(node.s)
    elif isinstance(node, ast.Subscript):
        return "%s[%s]" % (stringify(node.value), stringify(node.slice))
    elif isinstance(node, ast.Tuple):
        return ', '.join(stringify(e) for e in node.elts)
    else:
        # TODO: f-string
        raise NotImplementedError('Unable to render %s instance' % type(node).__name__)


class Signature:
    """A Signature object represents the call signature of a function and its return annotation.

    This class provides ``inspect.Signature`` like interface for string based call signatures::

        sig = Signature('(a, b, *args, **kwargs)')
        sig.parameters  # => returns a dict for Parameters
    """

    def __init__(self, signature: str) -> None:
        module = ast.parse('def func' + signature + ': pass')
        self.definition = cast(ast.FunctionDef, module.body[0])

    @cached_property
    def parameters(self) -> Mapping:
        args = self.definition.args
        params = []

        if hasattr(args, "posonlyargs"):  # for py38+
            for arg in args.posonlyargs:  # type: ignore
                annotation = stringify(arg.annotation) or Parameter.empty
                params.append(Parameter(arg.arg, Parameter.POSITIONAL_ONLY,
                                        annotation=annotation))

        for i, arg in enumerate(args.args):
            if len(args.args) - i <= len(args.defaults):
                default = stringify(args.defaults[-len(args.args) + i])
            else:
                default = Parameter.empty

            annotation = stringify(arg.annotation) or Parameter.empty
            params.append(Parameter(arg.arg, Parameter.POSITIONAL_OR_KEYWORD,
                                    default=default, annotation=annotation))

        if args.vararg:
            annotation = stringify(args.vararg.annotation) or Parameter.empty
            params.append(Parameter(args.vararg.arg, Parameter.VAR_POSITIONAL,
                                    annotation=annotation))

        for i, arg in enumerate(args.kwonlyargs):
            default = stringify(args.kw_defaults[i])
            annotation = stringify(arg.annotation) or Parameter.empty
            params.append(Parameter(arg.arg, Parameter.KEYWORD_ONLY, default=default,
                                    annotation=annotation))

        if args.kwarg:
            annotation = stringify(args.kwarg.annotation) or Parameter.empty
            params.append(Parameter(args.kwarg.arg, Parameter.VAR_KEYWORD,
                                    annotation=annotation))

        return OrderedDict((p.name, p) for p in params)

    @cached_property
    def return_annotation(self) -> str:
        return stringify(self.definition.returns)