# # This file is distributed under the MIT License. See LICENSE.md for details. # import ast import os from pathlib import Path from typing import List, cast import black from jinja2 import Environment from tuple_tree_generator.schema import Definition, EnumDefinition, ReferenceDefinition from tuple_tree_generator.schema import ScalarDefinition, Schema, SequenceDefinition from tuple_tree_generator.schema import StructDefinition, StructField, UpcastableDefinition from tuple_tree_generator.schema.struct import SimpleStructField from .jinja_utils import int_re, is_reference_struct_field, is_sequence_struct_field from .jinja_utils import is_simple_struct_field, loader class PythonGenerator: def __init__( self, schema: Schema, output, string_types=[], external_types=[], mixins_paths=[], ): self.schema = schema self.root_type = schema.root_type self.output = output self.string_types = string_types self.external_types = external_types self.mixins_paths = mixins_paths self.jinja_environment = Environment( block_start_string="##", block_end_string="##", variable_start_string="#{", variable_end_string="}#", comment_start_string="###", comment_end_string="###", loader=loader, ) self.jinja_environment.filters["python_type"] = self.python_type self.jinja_environment.filters["type_metadata"] = self.type_metadata self.jinja_environment.filters["docstring"] = self.render_docstring self.jinja_environment.filters["get_default_value"] = self.get_default_value self.jinja_environment.filters["gen_key"] = self.gen_key self.jinja_environment.filters["key_parser"] = self.key_parser self.jinja_environment.filters["get_mixins"] = self.get_mixins self.jinja_environment.globals["get_mixins_imports"] = self.get_mixins_imports self.jinja_environment.tests["simple_field"] = is_simple_struct_field self.jinja_environment.tests["sequence_field"] = is_sequence_struct_field self.jinja_environment.tests["reference_field"] = is_reference_struct_field self.template = self.jinja_environment.get_template("tuple_tree_gen.py.tpl") self.mixins: dict[str, List[str]] = self._parse_mixins() def emit_python(self) -> str: rendered_template = self.template.render( enums=self.schema.enum_definitions(), structs=self.schema.struct_definitions(), generator=self, schema=self.schema, root_type=self.schema.root_type, version=self.schema.version, ) return black.format_str( rendered_template, mode=black.Mode(line_length=100), ) @classmethod def python_type(cls, field: StructField | str): if isinstance(field, str): if field == "string": return "str" if int_re.match(field): return "int" else: return field resolved_type = field.resolved_type assert resolved_type return cls._python_type(resolved_type) @classmethod def _python_type(cls, resolved_type: Definition): assert isinstance(resolved_type, Definition) if isinstance(resolved_type, StructDefinition): return resolved_type.name if isinstance(resolved_type, SequenceDefinition): return f"TypedList[{cls._python_type(resolved_type.element_type)}]" if isinstance(resolved_type, EnumDefinition): return resolved_type.name if isinstance(resolved_type, ScalarDefinition): assert resolved_type.name if resolved_type.name == "string": return "str" if int_re.match(resolved_type.name): return "int" return resolved_type.name if isinstance(resolved_type, ReferenceDefinition): pointee = cls._python_type(resolved_type.pointee) root = cls._python_type(resolved_type.root) return f"Reference[{pointee}, {root}]" if isinstance(resolved_type, UpcastableDefinition): return cls._python_type(resolved_type.base) assert False def scalar_converter(self, type_name: str) -> str: if type_name == "string": return "str" if type_name == "bool": return "bool" if int_re.fullmatch(type_name): return "int" if type_name in self.external_types + self.string_types: return type_name raise Exception(f"Unexpected scalar: {type_name}") def _real_type(self, field: StructField) -> str: resolved_type = field.resolved_type if isinstance(resolved_type, SequenceDefinition): resolved_type = resolved_type.element_type if isinstance(resolved_type, UpcastableDefinition): resolved_type = resolved_type.base assert resolved_type real_type = resolved_type.name if isinstance(resolved_type, ScalarDefinition): return self.scalar_converter(real_type) return real_type def type_metadata(self, field: StructField) -> str: real_type = self._real_type(field) possible_values = None external = False if real_type in {"str", "bool", "int"}: hint_type = real_type ctor = "native" elif real_type in [e.name for e in self.schema.enum_definitions()]: hint_type = "str" ctor = "enum" possible_values = f"{real_type}" elif isinstance(field.resolved_type, ReferenceDefinition): hint_type = "Reference" ctor = "parse" elif real_type in [ *(s.name for s in self.schema.struct_definitions() if s.abstract), *(self.string_types), ]: hint_type = real_type ctor = "parse" if real_type in self.string_types: external = True else: hint_type = real_type ctor = "class" if real_type in self.external_types: external = True assert hint_type is_sequence = isinstance(field.resolved_type, SequenceDefinition) if isinstance(field.resolved_type, SequenceDefinition): is_abstract = isinstance(field.resolved_type.element_type, UpcastableDefinition) else: is_abstract = isinstance(field.resolved_type, UpcastableDefinition) return ( f'{{"type": {hint_type}, ' + (f'"possible_values": {possible_values},' if possible_values is not None else "") + f'"ctor": "{ctor}", ' + f'"is_key": {self.to_bool(field.is_key)}, ' + f'"is_array": {self.to_bool(is_sequence)}, ' + f'"is_abstract": {self.to_bool(is_abstract)},' + f'"external": {self.to_bool(external)}}}' ) @staticmethod def to_bool(_input: bool) -> str: return "True" if _input else "False" def _parse_mixins(self) -> dict[str, List[str]]: """Parse the mixins stored in `mixins_paths`. These files should define one or more of the following: * `AllMixin`: a mixin that will be applied to all the classes created by this generator * `{Name}Mixin`: this mixin will be applied to the class with `Name`, e.g. BinaryMixin will be applied to the class 'Binary' These classes can be defined either via class declaration or assignment, for example: ```python class AllMixin: ... class _CommonMixin: ... FooMixin = _CommonMixin BarMixin = _CommonMixin ``` """ mixins: dict[str, List[str]] = {} for mixins_path in self.mixins_paths: with open(Path(mixins_path), "r") as mixin_file: mixins_parsed = ast.parse(mixin_file.read()) mixin_classes = [] for node in mixins_parsed.body: # This node is encountered when defining a class, e.g. # class Foo: ... # noqa: E800 if isinstance(node, ast.ClassDef): mixin_classes.append(node.name) # This node is encountered when checking assignments e.g. # FooMixin = _CommonMixin # noqa: E800 elif isinstance(node, ast.Assign): for target in node.targets: mixin_classes.append(cast(ast.Name, target).id) mixins[mixins_path] = mixin_classes return mixins def get_mixins(self, struct_name: str) -> str: for mixin_classes in self.mixins.values(): for mc in mixin_classes: if struct_name == mc.removesuffix("Mixin"): return mc if "AllMixin" in mixin_classes: return "AllMixin" return "" def get_mixins_imports(self) -> List[str]: """ This method supports mixin imports from the `model` dir level or from subdirs """ imports: List[str] = [] for mixin_path, mixin_classes in self.mixins.items(): path_distance = os.path.relpath( os.path.dirname(mixin_path), os.path.dirname(self.output) ) file_name = os.path.basename(mixin_path).removesuffix(".py") if path_distance == ".": path = f".{file_name}" elif path_distance == "..": path = f"..{file_name}" else: # Transform from unix path to python import path = "" while path_distance.startswith("../"): path += "." path_distance = path_distance[3:] path += path_distance.replace("/", ".") path += f".{file_name}" imports.append(f"from {path} import {', '.join(mixin_classes)}") # noqa: E501 return imports @staticmethod def gen_key(struct: StructDefinition) -> str: fields = [f"{{self.{f.name}}}" for f in struct.key_fields] return f"f\"{'-'.join(fields)}\"" @staticmethod def key_parser(struct: StructDefinition) -> str: fields = [f'"{f.name}": parts[{i}]' for i, f in enumerate(struct.key_fields)] return ",".join(fields) @staticmethod def render_docstring(docstr: str, indent=1): if not docstr: return "" # Replace '\' with '\\' to avoid "SyntaxWarning: invalid escape sequence" lines = docstr.replace("\\", "\\\\").splitlines(keepends=False) rendered_docstring = '"""' rendered_docstring += ("\n" + " " * indent).join(lines) if len(lines) > 1: rendered_docstring += "\n " * indent rendered_docstring += '"""' return rendered_docstring def get_default_value(self, field: StructField): if isinstance(field.resolved_type, SequenceDefinition): return "[]" elif isinstance(field.resolved_type, ReferenceDefinition): return '""' elif isinstance(field, SimpleStructField): if field.type == "string" or field.type in self.string_types: assert not field.default or isinstance(field.default, str) return f'"{field.default if field.default else ""}"' elif field.type == "bool": assert not field.default or isinstance(field.default, bool) return "True" if field.default else "False" elif int_re.match(field.type): assert not field.default or isinstance(field.default, int) return f"{field.default if field.default else 0}" elif isinstance(field.resolved_type, EnumDefinition): assert not field.default or isinstance(field.default, str) return f'"{field.default if field.default else "Invalid"}"' assert not field.default, ( "Currently `default:` is only allowed on simple types: " "integers, booleans, strings and enums." ) if field.upcastable: return "None" else: return f"{field.type}()" else: raise ValueError()