regenerate_readme.py

import sys
import ast
import argparse
import inspect
from textwrap import dedent
from typing import Any
from contextlib import contextmanager
from string import whitespace
import re

import logging


logger = logging.getLogger(__name__)


class MyVisitor(ast.NodeVisitor):
    def __init__(self, source: str):
        super().__init__()
        self.source = source
        # Line number counting in ast starts at 1
        self.source_lines = [""] + source.splitlines(keepends=False)
        self.stack = []
        self.section_chars = '#*=-^"' + '"' * 1000
        self.flair_pattern = r"^\|[a-zA-Z_-]+\|$"
        self.all_filter = []

    @property
    def depth(self):
        return len(self.stack)

    @contextmanager
    def stacker(self, node):
        self.stack.append(node)
        try:
            yield
        finally:
            self.stack.pop()

    def make_titletext(self, text, prefix="", suffix="") -> str:
        text = prefix + text + suffix
        underline_len = len(text.encode())
        return "\n" + text + "\n" + self.section_chars[self.depth] * underline_len

    def dig(self, node):
        with self.stacker(node):
            super().generic_visit(node)

    def text_from_node_to_before_body(self, node) -> str:
        # text = ' '.join(self.source_lines[node.lineno:node.body[0].lineno])
        lines = self.source_lines[node.lineno : node.body[0].lineno]
        lines = map(str.strip, lines)
        lines = [re.sub(r",$", ", ", line) for line in lines]
        text = "".join(lines)
        return text.strip(whitespace + ":")

    def generic_visit(self, node: ast.AST) -> Any:
        logger.info(" " * 4 * len(self.stack) + "visited %s", node)
        self.dig(node)

    def visit_Module(self, node: ast.Module) -> Any:
        logger.info("visited module", node)
        self.dig(node)

    def stack_contains_a_function_before_a_class(self) -> False:
        for node in reversed(self.stack):
            if isinstance(node, ast.ClassDef):
                return False
            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                return True
        else:
            return False

    def scoop_flair_items_from_first_body_expr(self, node):
        # Extract flair tags out of docstring. Note: we don't print out the
        # docstring here, that will be left to the ``visit_Expr`` method.
        # We just want the flair so that the title can be modified
        docstring = ast.get_docstring(node)
        flair_items = []
        if docstring:
            flair_items = re.findall(self.flair_pattern, docstring, flags=re.MULTILINE)

        flair_text = " ".join(flair_items) + (" " if flair_items else "")
        return flair_text

    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
        if self.stack_contains_a_function_before_a_class():
            # No documentation for nested functions
            return

        if node.name.startswith("_"):
            return

        if self.depth == 1 and node.name not in self.all_filter:
            return

        text = self.text_from_node_to_before_body(node)
        text = re.sub(r"^def ", "", text)

        anchor_text = node.name
        if isinstance(self.stack[-1], ast.ClassDef):
            class_name = self.stack[-1].name
            text = f"{class_name}.{text}"
            anchor_text = f"{class_name}.{anchor_text}"

        decs = ""
        if node.decorator_list:
            logger.info(node.decorator_list)
            decs = (
                " ".join(
                    "@" + n.name if hasattr(n, "name") else "@" + n.id
                    for n in node.decorator_list
                )
                + " "
            )
        text = decs + text

        flair_text = self.scoop_flair_items_from_first_body_expr(node)
        text = flair_text + "``" + text + "``"

        print()
        print(f".. _{anchor_text}:")
        print()
        print(self.make_titletext(text))

        # Recurse
        self.dig(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
        logger.info("visited async function %s", node)
        self.visit_FunctionDef(node)
        self.dig(node)

    def visit_ClassDef(self, node: ast.ClassDef) -> Any:
        if node.name.startswith("_"):
            return

        if self.depth == 1 and node.name not in self.all_filter:
            return

        anchor_text = node.name
        print()
        print(f".. _{anchor_text}:")
        print()

        text = self.text_from_node_to_before_body(node)
        flair_text = self.scoop_flair_items_from_first_body_expr(node)
        text = flair_text + "``" + text + "``"
        print(self.make_titletext(text))
        self.dig(node)

    def visit_Expr(self, node: ast.Expr) -> Any:
        """Handles docstrings and any other strings that might just appear
        in various places."""
        is_constant_str = isinstance(node.value, ast.Constant) and isinstance(
            node.value.value, str
        )
        if not is_constant_str:
            return

        # First line might have different indentation!
        first_line, _, other_lines = node.value.value.partition("\n")
        text = first_line.strip() + "\n" + dedent(other_lines)
        text = re.sub(self.flair_pattern, "", text, flags=re.MULTILINE)
        # Replace doctest notifiers
        tokens = ["<BLANKLINE>", "# doctest: +SKIP", "# doctest: +ELLIPSIS"]
        for t in tokens:
            text = text.replace(t, " " * len(t))
        print(text)

    def visit_Assign(self, node: ast.Assign) -> Any:
        """This is only used to collect the __all__ values to filter
        the output for those."""
        if not hasattr(node.targets[0], "id"):
            return

        if node.targets[0].id != "__all__":
            return

        self.all_filter = [n.value for n in node.value.elts]


def main(args):
    file_contents = open(args.module).read()
    mod = ast.parse(file_contents, filename=args.module)

    visitor = MyVisitor(file_contents)
    visitor.visit(mod)


def make_titletext(name, element, level, prefix=""):
    section_chars = '#*=-^"'

    if inspect.isclass(element):
        text = f"class {name}"
    elif inspect.ismethod(element):
        # class methods
        sig = inspect.signature(element)
        text = f"*classmethod* ``{element.__qualname__}{sig}``"
    elif inspect.isfunction(element):
        # class methods
        sig = inspect.signature(element)
        text = f"``{element.__qualname__}{sig}``"
    else:
        try:
            sig = inspect.signature(element)
            text = f"``{element.__qualname__}{sig}``"
        except (ValueError, TypeError):
            text = name

    if prefix:
        text = prefix + text

    underline_len = len(text.encode())

    return "\n" + text + "\n" + section_chars[level] * underline_len


def parse(node, level=0, include_names=None):
    def elems_by_appearance(node):
        if node.__name__ == "type":
            return

        def line_number(node):
            try:
                return inspect.getsourcelines(node)[1]
            except TypeError:
                return 0

        def subnodes(node):
            for n in dir(node):
                try:
                    yield getattr(node, n)
                except AttributeError:
                    continue

        yield from sorted(subnodes(node), key=line_number)

    for element in elems_by_appearance(node):
        try:
            name = element.__name__
        except AttributeError:
            continue

        if name.startswith("_"):
            continue

        allow_emit = level != 1 or (
            not include_names or element.__name__ in include_names
        )
        if allow_emit:
            yield name, element, level
            if hasattr(element, "__doc__"):
                print(name, "hasdoc")
            if inspect.isclass(element):
                yield from parse(element, level + 1)


if __name__ == "__main__":
    logging.basicConfig(level="WARNING", stream=sys.stderr)
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--module", help="Module name to document")
    parser.add_argument("-a", "--all", help="Only document entries in __all__")
    args = parser.parse_args()
    main(args)