Class AddSchemasTransformer

scripts/openapi.py:1520–1566

class AddSchemasTransformer(CstTransformerBase):
    def __init__(self, class_name: str, schemas: list[str]):
        super().__init__()
        self.class_name = class_name
        self.schemas = schemas
        self.schema_added = 0

    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef):
        if self.current_class_name == self.class_name:
            docstring = get_class_docstring(updated_node)
            if not docstring:
                print("Class has no docstring")
            else:
                lines = docstring.splitlines()
                first_line = lines[1]
                indent = first_line[: len(first_line) - len(first_line.lstrip())]
                heading = len(lines) - 1  # if there is no heading, we place it before the last line (the closing """)
                empty_footing = lines[-2].strip() == ""
                schema_lines = []
                for idx, line in enumerate(lines):
                    if "The OpenAPI schema can be found at" in line:
                        heading = idx
                        schema_lines = lines[idx + 1 : -2 if empty_footing else -1]
                        break
                schema_lines = {schema_line for schema_line in schema_lines if schema_line}
                before = len(schema_lines)
                schema_lines = sorted(list(schema_lines.union({f"{indent}- {schema}" for schema in self.schemas})))
                after = len(schema_lines)
                lines = (
                    lines[:heading]
                    +
                    # we add an empty line before the schema lines title if there is none
                    ([""] if lines[heading - 1].strip() else [])
                    + [indent + "The OpenAPI schema can be found at"]
                    # we add an empty line after the schema lines title to get a proper bullet list in the docs
                    + [""]
                    + schema_lines
                    + [""]
                    + lines[-1:]
                )
                docstring = "\n".join(lines)
                stmt = cst.SimpleStatementLine([cst.Expr(cst.SimpleString(docstring))])
                stmts = [stmt] + list(updated_node.body.body[1:])
                updated_node = updated_node.with_changes(body=updated_node.body.with_changes(body=stmts))
                self.schema_added += after - before

        return super().leave_ClassDef(original_node, updated_node)


class CreateClassMethodTransformer(CstTransformerBase):