Source code for xsdata.formats.plantuml.generator

from pathlib import Path
from typing import Iterator
from typing import List
from typing import Tuple

from xsdata.generators import AbstractGenerator
from xsdata.models.codegen import Class
from xsdata.models.elements import Schema


class PlantUmlGenerator(AbstractGenerator):
    templates_dir = Path(__file__).parent.joinpath("templates")

    def render_module(self, output: str) -> str:
        return self.template("module").render(output=output)

    def render_class(self, obj: Class) -> str:
        template = "enum" if obj.is_enumeration else "class"
        return self.template(template).render(obj=obj)

    def render(
        self, schema: Schema, classes: List[Class], package: str
    ) -> Iterator[Tuple[Path, str]]:
        """Given a schema, a list of classes and a target package return to the
        writer factory the target file path and the rendered output."""
        module = schema.module
        package_arr = package.split(".")
        package = "{}.{}".format(".".join(package_arr), module)
        target = Path.cwd().joinpath(*package_arr)
        file_path = target.joinpath(f"{module}.pu")

        self.resolver.process(classes=classes, schema=schema, package=package)

        output = self.render_classes()

        yield file_path, self.render_module(output=output)

    def render_classes(self) -> str:
        """Sort classes by name and return the rendered output."""
        classes = sorted(self.resolver.sorted_classes(), key=lambda x: x.name)
        output = "\n".join(map(self.render_class, classes)).strip()
        return f"\n{output}\n"