|
| 1 | +"""Attribute docstrings parsing. |
| 2 | +
|
| 3 | +.. seealso:: https://peps.python.org/pep-0257/#what-is-a-docstring |
| 4 | +""" |
| 5 | +import ast |
| 6 | +import inspect |
| 7 | +import textwrap |
| 8 | +import typing as T |
| 9 | +from types import ModuleType |
| 10 | + |
| 11 | +from .common import Docstring |
| 12 | +from .common import DocstringParam |
| 13 | + |
| 14 | + |
| 15 | +def ast_get_constant_value(node: ast.AST) -> T.Any: |
| 16 | + """Return the constant's value if the given node is a constant.""" |
| 17 | + return getattr(node, 'value') |
| 18 | + |
| 19 | + |
| 20 | +def ast_unparse(node: ast.AST) -> T.Optional[str]: |
| 21 | + """Convert the AST node to source code as a string.""" |
| 22 | + if hasattr(ast, 'unparse'): |
| 23 | + return ast.unparse(node) |
| 24 | + # Support simple cases in Python < 3.9 |
| 25 | + if isinstance(node, ast.Constant): |
| 26 | + return str(ast_get_constant_value(node)) |
| 27 | + if isinstance(node, ast.Name): |
| 28 | + return node.id |
| 29 | + return None |
| 30 | + |
| 31 | + |
| 32 | +def ast_is_literal_str(node: ast.AST) -> bool: |
| 33 | + """Return True if the given node is a literal string.""" |
| 34 | + return ( |
| 35 | + isinstance(node, ast.Expr) |
| 36 | + and isinstance(node.value, ast.Constant) |
| 37 | + and isinstance(ast_get_constant_value(node.value), str) |
| 38 | + ) |
| 39 | + |
| 40 | + |
| 41 | +def ast_get_attribute( |
| 42 | + node: ast.AST, |
| 43 | +) -> T.Optional[T.Tuple[str, T.Optional[str], T.Optional[str]]]: |
| 44 | + """Return name, type and default if the given node is an attribute.""" |
| 45 | + if isinstance(node, (ast.Assign, ast.AnnAssign)): |
| 46 | + target = ( |
| 47 | + node.targets[0] if isinstance(node, ast.Assign) else node.target |
| 48 | + ) |
| 49 | + if isinstance(target, ast.Name): |
| 50 | + type_str = None |
| 51 | + if isinstance(node, ast.AnnAssign): |
| 52 | + type_str = ast_unparse(node.annotation) |
| 53 | + default = None |
| 54 | + if node.value: |
| 55 | + default = ast_unparse(node.value) |
| 56 | + return target.id, type_str, default |
| 57 | + return None |
| 58 | + |
| 59 | + |
| 60 | +class AttributeDocstrings(ast.NodeVisitor): |
| 61 | + """An ast.NodeVisitor that collects attribute docstrings.""" |
| 62 | + |
| 63 | + attr_docs: T.Dict[str, T.Tuple[str, T.Optional[str], T.Optional[str]]] = {} |
| 64 | + prev_attr = None |
| 65 | + |
| 66 | + def visit(self, node: T.Any) -> None: |
| 67 | + if self.prev_attr and ast_is_literal_str(node): |
| 68 | + attr_name, attr_type, attr_default = self.prev_attr |
| 69 | + self.attr_docs[attr_name] = ( |
| 70 | + ast_get_constant_value(node.value), |
| 71 | + attr_type, |
| 72 | + attr_default, |
| 73 | + ) |
| 74 | + self.prev_attr = ast_get_attribute(node) |
| 75 | + if isinstance(node, (ast.ClassDef, ast.Module)): |
| 76 | + self.generic_visit(node) |
| 77 | + |
| 78 | + def get_attr_docs( |
| 79 | + self, component: T.Any, |
| 80 | + ) -> T.Dict[str, T.Tuple[str, T.Optional[str], T.Optional[str]]]: |
| 81 | + """Get attribute docstrings from the given component. |
| 82 | +
|
| 83 | + :param component: component to process (class or module) |
| 84 | + :returns: for each attribute docstring, a tuple with (description, |
| 85 | + type, default) |
| 86 | + """ |
| 87 | + self.attr_docs = {} |
| 88 | + self.prev_attr = None |
| 89 | + try: |
| 90 | + source = textwrap.dedent(inspect.getsource(component)) |
| 91 | + except OSError: |
| 92 | + pass |
| 93 | + else: |
| 94 | + tree = ast.parse(source) |
| 95 | + if inspect.ismodule(component): |
| 96 | + self.visit(tree) |
| 97 | + elif isinstance(tree, ast.Module) and isinstance( |
| 98 | + tree.body[0], ast.ClassDef, |
| 99 | + ): |
| 100 | + self.visit(tree.body[0]) |
| 101 | + return self.attr_docs |
| 102 | + |
| 103 | + |
| 104 | +def add_attribute_docstrings( |
| 105 | + obj: T.Union[type, ModuleType], docstring: Docstring, |
| 106 | +) -> None: |
| 107 | + """Add attribute docstrings found in the object's source code. |
| 108 | +
|
| 109 | + :param obj: object from which to parse attribute docstrings |
| 110 | + :param docstring: Docstring object where found attributes are added |
| 111 | + :returns: list with names of added attributes |
| 112 | + """ |
| 113 | + params = set(p.arg_name for p in docstring.params) |
| 114 | + for arg_name, (description, type_name, default) in ( |
| 115 | + AttributeDocstrings().get_attr_docs(obj).items() |
| 116 | + ): |
| 117 | + if arg_name not in params: |
| 118 | + param = DocstringParam( |
| 119 | + args=['attribute', arg_name], |
| 120 | + description=description, |
| 121 | + arg_name=arg_name, |
| 122 | + type_name=type_name, |
| 123 | + is_optional=default is not None, |
| 124 | + default=default, |
| 125 | + ) |
| 126 | + docstring.meta.append(param) |
0 commit comments