clan-core/pkgs/classgen/main.py

import argparse
import json
from collections.abc import Callable, Iterable
from functools import partial
from pathlib import Path
from typing import Any


class Error(Exception):
    pass


def sort_types(items: Iterable[str]) -> list[str]:
    def sort_key(item: str) -> tuple[int, str]:
        # Priority order: lower number = higher priority
        if item.startswith(("dict", "list")):
            return (0, item)  # Highest priority, dicts and lists should be first
        if item == "None":
            return (2, item)  # Lowest priority, None should be last
        return (1, item)  # Middle priority, sorted alphabetically

    return sorted(items, key=sort_key)


# Function to map JSON schemas and types to Python types
def map_json_type(
    json_type: Any,
    nested_types: list[str] | None = None,
    parent: Any = None,
) -> list[str]:
    if nested_types is None:
        nested_types = ["Any"]
    if isinstance(json_type, list):
        res: list[str] = []
        for t in json_type:
            res.extend(map_json_type(t))
        return sort_types(set(res))

    if isinstance(json_type, dict):
        items = json_type.get("items")
        if items:
            nested_types = map_json_type(items)

        if not json_type.get("type") and json_type.get("tsType") == "unknown":
            return ["Unknown"]

        union = json_type.get("oneOf")
        if union:
            res: list[str] = []
            for t in union:
                res.extend(map_json_type(t, nested_types, parent))
            return sort_types(set(res))

        return sort_types(map_json_type(json_type.get("type"), nested_types))
    if json_type == "string":
        return ["str"]
    if json_type == "integer":
        return ["int"]
    if json_type == "number":
        return ["float"]
    if json_type == "boolean":
        return ["bool"]
    # In Python, "number" is analogous to the float type.
    # https://json-schema.org/understanding-json-schema/reference/numeric#number
    if json_type == "number":
        return ["float"]
    if json_type == "array":
        if not nested_types:
            msg = f"Array type not found for {parent}"
            raise Error(msg)
        return [f"""list[{" | ".join(sort_types(nested_types))}]"""]
    if json_type == "object":
        if not nested_types:
            msg = f"dict type not found for {parent}"
            raise Error(msg)
        return [f"""dict[str, {" | ".join(sort_types(nested_types))}]"""]
    if json_type == "null":
        return ["None"]

    msg = f"Python type not found for {json_type}"
    raise Error(msg)


known_classes = set()
# TODO: make this configurable
root_class = "Clan"


def field_def_from_default_type(
    field_name: str,
    field_types: list[str],
    class_name: str,
    finalize_field: Callable[..., tuple[str, str]],
) -> tuple[str, str] | None:
    if "dict" in str(field_types):
        return finalize_field(
            field_types=field_types,
            default_factory="dict",
        )

    if "list" in str(field_types):
        return finalize_field(
            field_types=field_types,
            default_factory="list",
        )
    if "None" in str(field_types):
        return finalize_field(
            field_types=field_types,
            default="None",
        )

    if class_name.endswith("Config"):
        # SingleDiskConfig
        # PackagesConfig
        # ...
        # Config classes MUST always be optional
        msg = f"""
            #################################################
            Clan module '{class_name}' specifies a top-level option '{field_name}' without a default value.

            To fix this:
            - Add a default value to the option

            lib.mkOption {{
                type = lib.types.nullOr lib.types.str;
                default = null; # <- Add a default value here
            }};

            # Other options

            - make the field nullable

            lib.mkOption {{
                #      ╔══════════════╗ <- Nullable type
                type = lib.types.nullOr lib.types.str;
            }};

            - Use lib.types.attrsOf if suitable
            - Use lib.types.listOf if suitable


            Or report this problem to the clan team. So the class generator can be improved.
            #################################################
            """
        raise Error(msg)

    return None


def field_def_from_default_value(
    default_value: Any,
    field_name: str,
    field_types: list[str],
    nested_class_name: str,
    finalize_field: Callable[..., tuple[str, str]],
) -> tuple[str, str] | None:
    # default_value = prop_info.get("default")
    if "Unknown" in field_types:
        # Unknown type, doesnt matter what the default value is
        # type Unknown | a -> Unknown
        return finalize_field(
            field_types=field_types,
        )
    if default_value is None:
        return finalize_field(
            field_types=[*field_types, "None"],
            default="None",
        )
    if isinstance(default_value, list):
        return finalize_field(
            field_types=field_types,
            default_factory="list",
        )
    if isinstance(default_value, dict):
        serialised_types = " | ".join(field_types)
        if serialised_types == nested_class_name:
            return finalize_field(
                field_types=field_types,
                default_factory=nested_class_name,
            )

        if "dict[str," in serialised_types:
            return finalize_field(
                field_types=field_types,
                default_factory="dict",
            )
        return finalize_field(
            field_types=field_types,
            default_factory="dict",
            type_appendix=" | dict[str,Any]",
        )

    # Primitive types
    if isinstance(default_value, str):
        return finalize_field(
            field_types=field_types,
            default=f"'{default_value}'",
        )
    if isinstance(default_value, bool | int | float):
        # Bool must be checked before int
        return finalize_field(
            field_types=field_types,
            default=f"{default_value}",
        )

    # Other default values unhandled yet.
    msg = f"Unhandled default value for field '{field_name}' - default value: {default_value}. ( In {nested_class_name} )"
    raise Error(msg)


def get_field_def(
    field_name: str,
    field_types: list[str],
    default: str | None = None,
    default_factory: str | None = None,
    type_appendix: str = "",
) -> tuple[str, str]:
    _field_types = set(field_types)
    if "None" in _field_types or default or default_factory:
        serialised_types = " | ".join(sort_types(_field_types)) + type_appendix
        serialised_types = f"{serialised_types}"
    else:
        serialised_types = " | ".join(sort_types(_field_types)) + type_appendix

    return (field_name, serialised_types)


# Recursive function to generate dataclasses from JSON schema
def generate_dataclass(
    schema: dict[str, Any],
    attr_path: list[str],
    class_name: str = root_class,
) -> str:
    properties = schema.get("properties", {})

    required_fields: list[tuple[str, str]] = []
    fields_with_default: list[tuple[str, str]] = []
    nested_classes: list[str] = []

    for prop, prop_info in properties.items():
        # If we are at the top level, and the attribute name is not explicitly included we only do shallow
        field_name = prop.replace("-", "_")

        # Skip "extraModules"
        # TODO: Introduce seperate model that is tied to the serialization format
        if "extraModules" in field_name:
            continue

        # if len(attr_path) == 0 and prop in shallow_attrs:
        #     field_def = field_name, "dict[str, Any]"
        #     fields_with_default.append(field_def)
        #     continue

        prop_type = prop_info.get("type", None)
        union_variants = prop_info.get("oneOf", [])
        enum_variants = prop_info.get("enum", [])

        # Collect all types
        field_types: list[str] = []

        title = prop_info.get("title", prop.removesuffix("s"))
        title_sanitized = "".join([p.capitalize() for p in title.split("-")])
        nested_class_name = f"""{class_name if class_name != root_class and not prop_info.get("title") else ""}{title_sanitized}"""

        if not prop_type and not union_variants and not enum_variants:
            prop_type = "Unknown"

        if union_variants:
            field_types = map_json_type(union_variants)

        elif prop_type == "array":
            item_schema = prop_info.get("items")

            if isinstance(item_schema, dict):
                field_types = map_json_type(
                    prop_type,
                    map_json_type(item_schema),
                    field_name,
                )
        elif enum := prop_info.get("enum"):
            literals = ", ".join([f'"{string}"' for string in enum])
            field_types = [f"""Literal[{literals}]"""]

        elif prop_type == "object":
            inner_type = prop_info.get("additionalProperties")
            if inner_type and inner_type.get("type") == "object":
                # Inner type is a class
                field_types = map_json_type(prop_type, [nested_class_name], field_name)

                if nested_class_name not in known_classes:
                    nested_classes.append(
                        generate_dataclass(
                            inner_type,
                            [*attr_path, prop],
                            nested_class_name,
                        ),
                    )
                    known_classes.add(nested_class_name)

            elif inner_type and inner_type.get("type") != "object":
                # Trivial type:
                # dict[str, inner_type]
                field_types = [
                    f"""dict[str, {" | ".join(map_json_type(inner_type))}]""",
                ]

            elif not inner_type:
                # The type is a class
                field_types = [nested_class_name]
                if nested_class_name not in known_classes:
                    nested_classes.append(
                        generate_dataclass(
                            prop_info,
                            [*attr_path, prop],
                            nested_class_name,
                        ),
                    )
                    known_classes.add(nested_class_name)
        elif prop_type == "Unknown":
            field_types = ["Unknown"]
        else:
            field_types = map_json_type(
                prop_type,
                nested_types=[],
                parent=field_name,
            )

        if not field_types:
            msg = f"Python type not found for {prop} {prop_info}"
            raise Error(msg)

        finalize_field = partial(get_field_def, field_name)

        # Sort and remove potential duplicates
        field_types_sorted = sort_types(set(field_types))

        if "default" in prop_info or field_name not in prop_info.get("required", []):
            if prop_info.get("type") == "object":
                prop_info.update({"default": {}})

            if "default" in prop_info:
                default_value = prop_info.get("default")
                field_def = field_def_from_default_value(
                    default_value=default_value,
                    field_name=field_name,
                    field_types=field_types_sorted,
                    nested_class_name=nested_class_name,
                    finalize_field=finalize_field,
                )
                if field_def:
                    fields_with_default.append(field_def)

                if not field_def:
                    # Finalize without the default value
                    field_def = finalize_field(
                        field_types=field_types_sorted,
                    )
                    required_fields.append(field_def)

            if "default" not in prop_info:
                # Field is not required and but also specifies no default value
                # Trying to infer default value from type
                field_def = field_def_from_default_type(
                    field_name=field_name,
                    field_types=field_types_sorted,
                    class_name=class_name,
                    finalize_field=finalize_field,
                )

                if field_def:
                    fields_with_default.append(field_def)
                if not field_def:
                    field_def = finalize_field(
                        field_types=field_types_sorted,
                    )
                    required_fields.append(field_def)

        else:
            field_def = finalize_field(
                field_types=field_types_sorted,
            )
            required_fields.append(field_def)

    # Join field name with type to form a complete field declaration
    # e.g. "name: str"
    all_field_declarations = [f"{n}: {t}" for n, t in (required_fields)] + [
        f"{n}: NotRequired[{class_name}{n.capitalize()}Type]"
        for n, t in (fields_with_default)
    ]
    hoisted_types: str = "\n".join(
        [
            f"{class_name}{n.capitalize()}Type = {x}"
            for n, x in (required_fields + fields_with_default)
        ],
    )
    fields_str = "\n    ".join(all_field_declarations)
    nested_classes_str = "\n\n".join(nested_classes)

    class_def = f"\n\n{hoisted_types}\n"
    class_def += f"\nclass {class_name}(TypedDict):\n"
    if not required_fields + fields_with_default:
        class_def += "    pass"
    else:
        class_def += f"    {fields_str}"

    return f"{nested_classes_str}\n\n{class_def}" if nested_classes_str else class_def


def run_gen(args: argparse.Namespace) -> None:
    print(f"Converting {args.input} to {args.output}")

    dataclass_code = ""
    with args.input.open() as f:
        schema = json.load(f)
        dataclass_code = generate_dataclass(schema, [])

    with args.output.open("w") as f:
        f.write(
            """# DO NOT EDIT THIS FILE MANUALLY. IT IS GENERATED.
# This file was generated by running `pkgs/clan-cli/clan_lib.inventory/update.sh`
#
# fmt: off
from typing import Any, Literal, NotRequired, TypedDict\n

# Mimic "unknown".
# 'Any' is unsafe because it allows any operations
# This forces the user to use type-narrowing or casting in the code
class Unknown:
    pass
""",
        )
        f.write(dataclass_code)
        f.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", help="Input JSON schema file", type=Path)
    parser.add_argument("output", help="Output Python file", type=Path)
    parser.add_argument(
        "--stop-at",
        type=str,
        help="Property name to stop generating classes for. Other classes below that property will be generated",
        default=None,
    )
    parser.set_defaults(func=run_gen)
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()