preserve_defaults.py

"""Preserve function defaults.

Preserve the default argument values of function signatures in source code
and keep them not evaluated for readability.
"""

import ast
import inspect
import sys
from inspect import Parameter
from typing import Any, Dict, List, Optional

from sphinx.application import Sphinx
from sphinx.locale import __
from sphinx.pycode.ast import parse as ast_parse
from sphinx.pycode.ast import unparse as ast_unparse
from sphinx.util import logging

logger = logging.getLogger(__name__)


class DefaultValue:
    def __init__(self, name: str) -> None:
        self.name = name

    def __repr__(self) -> str:
        return self.name


def get_function_def(obj: Any) -> ast.FunctionDef:
    """Get FunctionDef object from living object.
    This tries to parse original code for living object and returns
    AST node for given *obj*.
    """
    try:
        source = inspect.getsource(obj)
        if source.startswith((' ', r'\t')):
            # subject is placed inside class or block.  To read its docstring,
            # this adds if-block before the declaration.
            module = ast_parse('if True:\n' + source)
            return module.body[0].body[0]  # type: ignore
        else:
            module = ast_parse(source)
            return module.body[0]  # type: ignore
    except (OSError, TypeError):  # failed to load source code
        return None


def get_default_value(lines: List[str], position: ast.AST) -> Optional[str]:
    try:
        if sys.version_info < (3, 8):  # only for py38+
            return None
        elif position.lineno == position.end_lineno:
            line = lines[position.lineno - 1]
            return line[position.col_offset:position.end_col_offset]
        else:
            # multiline value is not supported now
            return None
    except (AttributeError, IndexError):
        return None


def update_defvalue(app: Sphinx, obj: Any, bound_method: bool) -> None:
    """Update defvalue info of *obj* using type_comments."""
    if not app.config.autodoc_preserve_defaults:
        return

    try:
        lines = inspect.getsource(obj).splitlines()
        if lines[0].startswith((' ', r'\t')):
            lines.insert(0, '')  # insert a dummy line to follow what get_function_def() does.
    except (OSError, TypeError):
        lines = []

    try:
        function = get_function_def(obj)
        if function.args.defaults or function.args.kw_defaults:
            sig = inspect.signature(obj)
            defaults = list(function.args.defaults)
            kw_defaults = list(function.args.kw_defaults)
            parameters = list(sig.parameters.values())
            for i, param in enumerate(parameters):
                if param.default is param.empty:
                    if param.kind == param.KEYWORD_ONLY:
                        # Consume kw_defaults for kwonly args
                        kw_defaults.pop(0)
                else:
                    if param.kind in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
                        default = defaults.pop(0)
                        value = get_default_value(lines, default)
                        if value is None:
                            value = ast_unparse(default)  # type: ignore
                        parameters[i] = param.replace(default=DefaultValue(value))
                    else:
                        default = kw_defaults.pop(0)
                        value = get_default_value(lines, default)
                        if value is None:
                            value = ast_unparse(default)  # type: ignore
                        parameters[i] = param.replace(default=DefaultValue(value))

            if bound_method and inspect.ismethod(obj):
                # classmethods
                cls = inspect.Parameter('cls', Parameter.POSITIONAL_OR_KEYWORD)
                parameters.insert(0, cls)

            sig = sig.replace(parameters=parameters)
            if bound_method and inspect.ismethod(obj):
                # classmethods can't be assigned __signature__ attribute.
                obj.__dict__['__signature__'] = sig
            else:
                obj.__signature__ = sig
    except (AttributeError, TypeError):
        # failed to update signature (ex. built-in or extension types)
        pass
    except NotImplementedError as exc:  # failed to ast.unparse()
        logger.warning(__("Failed to parse a default argument value for %r: %s"), obj, exc)


def setup(app: Sphinx) -> Dict[str, Any]:
    app.add_config_value('autodoc_preserve_defaults', False, True)
    app.connect('autodoc-before-process-signature', update_defvalue)

    return {
        'version': '1.0',
        'parallel_read_safe': True
    }