atlas-init 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

atlas_init/tf_ext/schema_to_dataclass.py

@@ -0,0 +1,522 @@
+from __future__ import annotations
+
+import keyword
+import logging
+import re
+from collections import defaultdict
+from dataclasses import fields
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any, Callable, ClassVar, Self
+
+from ask_shell import ShellError, run_and_wait
+from inflection import singularize
+from model_lib import Entity
+from pydantic import model_validator
+from zero_3rdparty import humps
+from zero_3rdparty.file_utils import copy, update_between_markers
+
+from atlas_init.tf_ext.models_module import (
+    ModuleGenConfig,
+    ResourceAbs,
+    ResourceTypePythonModule,
+    import_resource_type_python_module,
+)
+from atlas_init.tf_ext.provider_schema import ResourceSchema, SchemaAttribute, SchemaBlock
+from atlas_init.tf_ext.py_gen import (
+    as_set,
+    dataclass_matches,
+    ensure_dataclass_use_conversion,
+    import_from_path,
+    longest_common_substring_among_all,
+    make_post_init_line_optional,
+    module_dataclasses,
+    move_main_call_to_end,
+    primitive_types,
+)
+
+logger = logging.getLogger(__name__)
+
+MARKER_START = "# codegen atlas-init-marker-start"
+MARKER_END = "# codegen atlas-init-marker-end"
+
+
+def is_computed_only(attr: SchemaAttribute) -> bool:
+    return bool(attr.computed) and not bool(attr.required) and not bool(attr.optional)
+
+
+def type_from_schema_attr(attr: SchemaAttribute, parent_class_name=None, attr_name=None) -> str:
+    # Only handle attribute types (not nested_type)
+    t = attr.type
+    if isinstance(t, str):
+        return {
+            "string": "str",
+            "number": "float",
+            "bool": "bool",
+            "int": "int",
+            "any": "Any",
+        }.get(t, "Any")
+    elif isinstance(t, list):
+        # Terraform types: ["list", "string"] or ["set", "object", {...}]
+        if t[0] in ("list", "set"):
+            if len(t) == 2 and isinstance(t[1], str):
+                return f"List[{type_from_schema_attr(SchemaAttribute(type=t[1]))}]"
+            elif len(t) == 3 and isinstance(t[2], dict):
+                # object type
+                return "List[dict]"
+        elif t[0] == "map":
+            return "Dict[str, Any]"
+    elif isinstance(t, dict):
+        return "dict"
+    return "Any"
+
+
+def safe_name(name):
+    return f"{name}_" if keyword.iskeyword(name) else name
+
+
+def py_type_from_element_type(elem_type_val: str | dict[str, str] | Any) -> str:
+    if isinstance(elem_type_val, str):
+        return {
+            "string": "str",
+            "number": "float",
+            "bool": "bool",
+            "int": "int",
+            "any": "Any",
+        }.get(elem_type_val, "Any")
+    elif isinstance(elem_type_val, dict):
+        return "dict"
+    else:
+        return "Any"
+
+
+class DcField(Entity):
+    NO_DEFAULT: ClassVar[str] = "None"
+    METADATA_DEFAULT_NAME: ClassVar[str] = "default_hcl"
+    name: str
+    type_annotation: str
+    description: str | None = None
+    default_value: str = NO_DEFAULT
+    default_hcl_string: str | None = None
+    nested_class_name: str = ""
+    required: bool = False
+    optional: bool = False
+    computed: bool = False
+
+    @model_validator(mode="after")
+    def validate_self(self) -> Self:
+        self.name = safe_name(self.name)
+        return self
+
+    @property
+    def is_list(self) -> bool:
+        return self.type_annotation.startswith("List[")
+
+    @property
+    def is_dict(self) -> bool:
+        return self.type_annotation.startswith("Dict[")
+
+    @property
+    def is_nested(self) -> bool:
+        return self.type_annotation.startswith(("List[", "Dict[", "Set[")) or self.nested_class_name != ""
+
+    @property
+    def metadata(self) -> dict:
+        return {
+            key: value
+            for key, value in [
+                ("description", self.description),
+                (self.METADATA_DEFAULT_NAME, self.default_hcl_string),
+            ]
+            if value
+        }
+
+    @property
+    def declare(self) -> str:
+        if metadata := self.metadata:
+            field_args = ["default=None", f"metadata={metadata}"]
+            return f"    {self.name}: Optional[{self.type_annotation}] = field({', '.join(field_args)})"
+        return f"    {self.name}: Optional[{self.type_annotation}] = None"
+
+    @property
+    def declare_required(self) -> str:
+        """Why not use self.required? Even though an attribute is required in the schema, we might be able to infer the value, for example cluster_type"""
+        if metadata := self.metadata:
+            field_args = [f"metadata={metadata}"]
+            return f"    {self.name}: {self.type_annotation} = field({', '.join(field_args)})"
+        return f"    {self.name}: {self.type_annotation}"
+
+    @property
+    def post_init(self) -> str:
+        if cls_name := self.nested_class_name:
+            return make_post_init_line_optional(self.name, cls_name, is_list=self.is_list, is_map=self.is_dict)
+        return ""
+
+    @property
+    def computed_only(self) -> bool:
+        return self.computed and not self.required and not self.optional
+
+
+def nested_type_annotation(elem_cls_name: str, nesting_mode: str | None) -> str:
+    nesting_mode = nesting_mode or ""
+    if nesting_mode == "list":
+        return f"List[{elem_cls_name}]"
+    if nesting_mode == "set":
+        return f"Set[{elem_cls_name}]"
+    return elem_cls_name
+
+
+def convert_to_dataclass(
+    schema: ResourceSchema, existing: ResourceTypePythonModule, config: ModuleGenConfig, resource_type: str
+) -> str:
+    class_defs = []
+
+    def block_to_class(block: SchemaBlock, class_name: str, extra_post_init: list[str] | None = None) -> str:
+        lines = ["@dataclass", f"class {class_name}:"]
+        dc_fields: list[DcField] = []
+
+        def add_attribute(
+            attr_name: str,
+            attr: SchemaAttribute,
+            type_annotation: str,
+            *,
+            required: bool = False,
+            optional: bool = False,
+            computed: bool = False,
+            nested_class_name: str = "",
+        ):
+            dc_field = DcField(
+                name=attr_name,
+                type_annotation=type_annotation,
+                required=required or bool(attr.required),
+                optional=optional or bool(attr.optional),
+                computed=computed or bool(attr.computed),
+                description=attr.description,
+                nested_class_name=nested_class_name,
+            )
+            dc_fields.append(dc_field)
+
+        def add_block_attribute(
+            attr_name: str,
+            block_type: SchemaBlock,
+            required: bool = False,
+            optional: bool = False,
+            computed: bool = False,
+            description: str | None = None,
+        ):
+            nested_class_name = f"{class_name}_{attr_name.capitalize()}"
+            type_annotation = nested_type_annotation(nested_class_name, block_type.nesting_mode)
+            dc_field = DcField(
+                name=attr_name,
+                type_annotation=type_annotation,
+                required=required,
+                optional=optional,
+                computed=computed,
+                description=description,
+                nested_class_name=nested_class_name,
+            )
+            class_defs.append(block_to_class(block_type, nested_class_name))
+            dc_fields.append(dc_field)
+
+        for attr_name, attr in (block.attributes or {}).items():
+            if attr.deprecated or attr.deprecated_message or attr_name in config.skip_variables_extra(resource_type):
+                if attr.deprecated:
+                    logger.info(f"skipping deprecated attribute {attr_name} for {resource_type}")
+                continue
+            required = bool(attr.required)
+            if nested_block := attr.nested_type:
+                add_block_attribute(
+                    attr_name,
+                    nested_block,
+                    required=required,
+                    optional=bool(attr.optional),
+                    computed=bool(attr.computed),
+                    description=attr.description,
+                )
+                continue
+            if elem_type_val := attr.element_type:
+                elem_py_type = py_type_from_element_type(elem_type_val)
+                match attr.type:
+                    case ["map", *_]:
+                        py_type = f"Dict[str, {elem_py_type}]"
+                    case ["list", *_]:
+                        py_type = f"List[{elem_py_type}]"
+                    case ["set", *_]:
+                        py_type = f"Set[{elem_py_type}]"
+                    case _:
+                        py_type = elem_py_type
+                nested_class_name = ""
+                if elem_py_type not in ("str", "float", "bool", "int", "Any", "dict"):
+                    nested_class_name = elem_py_type
+                add_attribute(attr_name, attr, py_type, nested_class_name=nested_class_name)
+            else:
+                py_type = type_from_schema_attr(attr, class_name, attr_name)
+                add_attribute(attr_name, attr, py_type)
+
+        block_attributes = set()
+        for block_type_name, block_type in (block.block_types or {}).items():
+            if block_type.deprecated or block_type_name in config.skip_variables_extra(resource_type):
+                logger.info(f"skipping deprecated block type {block_type_name}")
+                continue
+            is_required = (block_type.min_items or 0) > 0 or bool(block_type.required)
+            block_attributes.add(block_type_name)
+            add_block_attribute(
+                block_type_name,
+                block_type.block_with_nesting_mode,
+                required=is_required,
+                optional=bool(block_type.optional),
+                description=block_type.description,
+            )
+
+        lines.append(
+            f"    {ResourceAbs.BLOCK_ATTRIBUTES_NAME}: ClassVar[Set[str]] = {as_set(sorted(block_attributes))}"
+        )
+        lines.append(
+            f"    {ResourceAbs.NESTED_ATTRIBUTES_NAME}: ClassVar[Set[str]] = {as_set([dc_field.name for dc_field in dc_fields if dc_field.is_nested])}"
+        )
+        lines.append(
+            f"    {ResourceAbs.REQUIRED_ATTRIBUTES_NAME}: ClassVar[Set[str]] = {as_set([dc_field.name for dc_field in dc_fields if dc_field.required])}"
+        )
+        lines.append(
+            f"    {ResourceAbs.COMPUTED_ONLY_ATTRIBUTES_NAME}: ClassVar[Set[str]] = {as_set([dc_field.name for dc_field in dc_fields if dc_field.computed_only])}"
+        )
+        default_strings = {
+            dc_field.name: default_hcl_string
+            for dc_field in dc_fields
+            if (default_hcl_string := config.attribute_default_hcl_strings(resource_type).get(dc_field.name))
+        }
+        lines.append(f"    {ResourceAbs.DEFAULTS_HCL_STRINGS_NAME}: ClassVar[dict[str, str]] = {default_strings!r}")
+
+        if not dc_fields:
+            lines.append("    pass")
+            return "\n".join(lines)
+        if not config.use_descriptions:
+            for dc_field in dc_fields:
+                dc_field.description = None
+        required_vars = config.required_variables(resource_type)
+        lines.extend(dc_field.declare_required for dc_field in dc_fields if dc_field.name in required_vars)
+        lines.extend(dc_field.declare for dc_field in dc_fields if dc_field.name not in required_vars)
+        post_init_lines = [post_init for dc_field in dc_fields if (post_init := dc_field.post_init)]
+        if extra_post_init:
+            post_init_lines.extend(extra_post_init)
+        if post_init_lines:
+            lines.append("    def __post_init__(self):")
+            lines.extend(post_init_lines)
+        lines.extend(["", ""])
+        return "\n".join(lines)
+
+    root_class_name = "Resource"
+    class_defs.append(block_to_class(schema.block, root_class_name, existing.extra_post_init_lines))
+
+    import_lines = [
+        "import json",
+        "import sys",
+        "from dataclasses import asdict, dataclass, field",
+        "from typing import Optional, List, Dict, Any, Set, ClassVar, Union, Iterable",
+    ]
+    import_lines.extend(existing.extra_import_lines)
+
+    module_str = "\n".join(import_lines + [""] + class_defs)
+    return module_str.strip() + "\n"
+
+
+_primitive_conversion = """
+    
+def format_primitive(value: Union[str, float, bool, int, None]):
+    if value is None:
+        return None
+    if value is True:
+        return "true"
+    if value is False:
+        return "false"
+    return str(value)
+"""
+_debug_logs = """
+    from pathlib import Path
+    logs_out = Path(__file__).parent / "logs.json"
+    logs_out.write_text(json_str)
+"""
+
+
+def main_entrypoint(existing: ResourceTypePythonModule, config: ModuleGenConfig) -> str:
+    parse_cls = "ResourceExt" if existing.resource_ext_cls_used else "Resource"
+    errors_func_call = r'"\\n".join(errors(resource))' if existing.errors_func_used else '""'
+    modify_out_func_call = "\n    resource = modify_out(resource)" if existing.modify_out_func_used else ""
+    logs_debug = _debug_logs if config.debug_json_logs else ""
+    return (
+        _primitive_conversion
+        + f"""
+def main():
+    input_data = sys.stdin.read()
+    # Parse the input as JSON
+    params = json.loads(input_data)
+    input_json = params["input_json"]
+    resource = {parse_cls}(**json.loads(input_json))
+    error_message = {errors_func_call}
+    primitive_types = ({", ".join(t.__name__ for t in primitive_types)}){modify_out_func_call}
+    output = {{
+        key: format_primitive(value) if value is None or isinstance(value, primitive_types) else json.dumps(value)
+        for key, value in asdict(resource).items()
+    }}
+    output["error_message"] = error_message
+    json_str = json.dumps(output){logs_debug}
+    print(json_str)
+if __name__ == "__main__":
+    main()
+
+"""
+    )
+
+
+def py_file_validate_and_auto_fixes(code: str, error_hint: str = "") -> str:
+    with TemporaryDirectory() as tmp_dir:
+        tmp_file = Path(tmp_dir) / "dataclass.py"
+        tmp_file.write_text(code)
+        run_fmt_and_fixes(tmp_file)
+        return tmp_file.read_text()
+
+
+def run_fmt_and_fixes(file_path: Path, error_hint: str = ""):
+    tmp_dir = file_path.parent
+    try:
+        run_and_wait("ruff format . --line-length 120", cwd=tmp_dir)
+    except ShellError as e:
+        logger.exception(f"Failed to format dataclass:\n{file_path.read_text()}\n{error_hint}")
+        raise e
+    try:
+        run_and_wait("ruff check --fix .", cwd=tmp_dir)
+    except ShellError as e:
+        logger.exception(f"Failed to check dataclass:\n{file_path.read_text()}\n{error_hint}")
+        raise e
+    return file_path.read_text()
+
+
+def dataclass_id(cls: type) -> str:
+    field_names = ",".join(sorted(f.name for f in fields(cls)))
+    computed_only_names = ",".join(sorted(f.name for f in fields(cls) if ResourceAbs.is_computed_only(f.name, cls)))
+    required_only_names = ",".join(sorted(f.name for f in fields(cls) if ResourceAbs.is_required(f.name, cls)))
+    id_parts = [field_names]
+    if computed_only_names:
+        id_parts.append(f"computed={computed_only_names}")
+    if required_only_names:
+        id_parts.append(f"required={required_only_names}")
+    return "|".join(id_parts)
+
+
+class NameAlreadyTakenError(Exception):
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)
+
+
+def simplify_classes(py_code: str) -> tuple[str, set[str]]:
+    with TemporaryDirectory() as tmp_dir:
+        tmp_file = Path(tmp_dir) / "dataclass.py"
+        tmp_file.write_text(py_code)
+        module = import_from_path("dataclass", tmp_file)
+        dataclasses = module_dataclasses(module)
+        fields_to_dataclass = defaultdict(list)
+        for name, dc in dataclasses.items():
+            fields_to_dataclass[dataclass_id(dc)].append(name)
+        new_names: set[str] = set()
+
+        def add_new_name(new_name: str) -> None:
+            if new_name in new_names or new_name in dataclasses:
+                raise NameAlreadyTakenError(f"Duplicate new name: {new_name}")
+            new_names.add(new_name)
+
+        if duplicate_classes := {k: v for k, v in fields_to_dataclass.items() if len(v) > 1}:
+            for duplicates in duplicate_classes.values():
+                py_code, new_name = rename_and_remove_duplicates(duplicates, py_code, add_new_name)
+        for old_classes in fields_to_dataclass.values():
+            if len(old_classes) != 1:
+                continue
+            cls_name = old_classes[0]
+            if "_" not in cls_name:
+                continue
+            new_name = extract_last_name_part(cls_name)
+            py_code = _safe_replace(py_code, cls_name, new_name)
+            add_new_name(new_name)
+        return py_code, new_names
+
+
+def _safe_replace(text: str, old: str, new: str) -> str:
+    def replacer(match: re.Match) -> str:
+        return match[0].replace(old, new)
+
+    return re.sub(rf"\W({old})\W", replacer, text)
+
+
+_plural_exception_list = {"Aws"}
+_cls_exception_mapping = {
+    "Resource": "ResourceElem",
+}
+
+
+def extract_last_name_part(full_name: str) -> str:
+    included_words = []
+    for word in reversed(full_name.split("_")):
+        included_words.insert(0, word)
+        if word[0].isupper():
+            break
+    plural_word = humps.pascalize("_".join(included_words))
+    name = plural_word
+    if plural_word not in _plural_exception_list:
+        name = singularize(plural_word)
+    return _cls_exception_mapping.get(name, name)
+
+
+def rename_and_remove_duplicates(
+    duplicates: list[str], py_code: str, add_new_name: Callable[[str], None]
+) -> tuple[str, str]:
+    duplicates_short = [extract_last_name_part(d) for d in duplicates]
+    new_name = longest_common_substring_among_all(duplicates_short)
+    try:
+        add_new_name(new_name)
+    except NameAlreadyTakenError:
+        new_name += "2"
+        add_new_name(new_name)
+    for replace in duplicates:
+        py_code = _safe_replace(py_code, replace, new_name)
+    py_code = remove_duplicates(py_code, new_name)
+    return py_code, new_name
+
+
+def remove_duplicates(py_code: str, new_name) -> str:
+    matches = list(dataclass_matches(py_code, new_name))
+    logger.info(f"found {len(matches)} matches for {new_name}")
+    while len(matches) > 1:
+        next_match = matches.pop()
+        py_code = py_code[: next_match.index_start] + py_code[next_match.index_end :]
+    return py_code
+
+
+SKIP_FILTER = {"Resource", "ResourceExt"}
+
+
+def generate_python_from_schema(
+    py_module: ResourceTypePythonModule, schema: ResourceSchema, config: ModuleGenConfig, resource_type: str
+) -> str:
+    dataclass_unformatted = convert_to_dataclass(schema, py_module, config, resource_type)
+    dataclass_unformatted = simplify_classes(dataclass_unformatted)[0]
+    return f"{dataclass_unformatted}\n{main_entrypoint(py_module, config)}"
+
+
+def convert_and_format(
+    resource_type: str,
+    schema: ResourceSchema,
+    config: ModuleGenConfig,
+    existing_path: Path | None = None,
+) -> str:
+    if existing_path is not None and existing_path.exists():
+        py_module = import_resource_type_python_module(resource_type, existing_path)
+        with TemporaryDirectory() as tmp_path:
+            tmp_file = Path(tmp_path) / f"{resource_type}.py"
+            copy(existing_path, tmp_file)
+            dataclass_unformatted = generate_python_from_schema(py_module, schema, config, resource_type)
+            update_between_markers(tmp_file, dataclass_unformatted, MARKER_START, MARKER_END)
+            move_main_call_to_end(tmp_file)
+            ensure_dataclass_use_conversion(py_module.dataclasses, tmp_file, SKIP_FILTER)
+            return run_fmt_and_fixes(tmp_file)
+    existing = ResourceTypePythonModule(resource_type)
+    dataclass_unformatted = generate_python_from_schema(existing, schema, config, resource_type)
+    return py_file_validate_and_auto_fixes(f"{MARKER_START}\n{dataclass_unformatted}\n{MARKER_END}\n")