Class ApplySchemaTransformer

scripts/openapi.py:840–1259

class ApplySchemaTransformer(ApplySchemaBaseTransformer):
    def __init__(
        self,
        module_name: str,
        class_name: str,
        properties: dict[str, (PythonType | GithubClass | None, bool)],
        completable: bool,
        deprecate: bool,
    ):
        super().__init__(module_name, class_name, properties, deprecate)
        self.completable = completable

    @staticmethod
    def deprecate_function(node: cst.FunctionDef) -> cst.FunctionDef:
        decorators = list(node.decorators)
        decorators.append(cst.Decorator(decorator=cst.Name(value="deprecated")))
        return node.with_changes(decorators=decorators)

    def inner_github_type(self, data_type: PythonType | GithubClass | list[PythonType | GithubClass]) -> [GithubClass]:
        if data_type is None:
            return []
        if isinstance(data_type, list):
            return [ght for dt in data_type for ght in self.inner_github_type(dt)]
        if isinstance(data_type, PythonType):
            return self.inner_github_type(data_type.inner_types)
        if isinstance(data_type, GithubClass):
            return [data_type]
        raise ValueError("Unsupported data type", data_type)

    def leave_Module(self, original_node: Module, updated_node: Module) -> Module:
        i = 0
        node = updated_node

        # add from __future__ import annotations if not the first import
        node = self.add_future_import(node)

        property_classes = {
            ghc
            for p in self.all_properties
            for ghc in self.inner_github_type(p.data_type)
            if ghc.module != self.module_name and ghc.name != self.class_name
        }
        import_classes = sorted(property_classes, key=lambda c: c.module)
        typing_classes = sorted(property_classes, key=lambda c: c.module)
        # TODO: do not import this file itself
        datetime_exists = False
        in_github_imports = False
        needs_datetime_import = any(
            p.data_type.type == "datetime" for p in self.all_properties if isinstance(p.data_type, PythonType)
        )

        # insert import classes if needed
        while (
            i < len(node.body)
            and isinstance(node.body[i], cst.SimpleStatementLine)
            and isinstance(node.body[i].body[0], (cst.Import, cst.ImportFrom))
        ):
            if self.is_datetime_import(node.body[i].body[0]):