Add support for template in protocols

Author: dgarros

infrahub_sdk/protocols_generator/generator.py

@@ -2,7 +2,6 @@
 
 from collections.abc import Mapping
 from pathlib import Path
-from typing import Any
 
 import jinja2
 
@@ -16,6 +15,7 @@
     NodeSchemaAPI,
     ProfileSchemaAPI,
     RelationshipSchemaAPI,
+    TemplateSchemaAPI,
 )
 from .constants import ATTRIBUTE_KIND_MAP, TEMPLATE_FILE_NAME
 
@@ -30,6 +30,7 @@ def __init__(self, schema: dict[str, MainSchemaTypesAll]):
         self.generics: dict[str, GenericSchemaAPI | GenericSchema] = {}
         self.nodes: dict[str, NodeSchemaAPI | NodeSchema] = {}
         self.profiles: dict[str, ProfileSchemaAPI] = {}
+        self.templates: dict[str, TemplateSchemaAPI] = {}
 
         for name, schema_type in schema.items():
             if isinstance(schema_type, (GenericSchemaAPI, GenericSchema)):
@@ -38,6 +39,8 @@ def __init__(self, schema: dict[str, MainSchemaTypesAll]):
                 self.nodes[name] = schema_type
             if isinstance(schema_type, ProfileSchemaAPI):
                 self.profiles[name] = schema_type
+            if isinstance(schema_type, TemplateSchemaAPI):
+                self.templates[name] = schema_type
 
         self.base_protocols = [
             e
@@ -53,30 +56,44 @@ def __init__(self, schema: dict[str, MainSchemaTypesAll]):
         self.sorted_profiles = self._sort_and_filter_models(
             self.profiles, filters=["CoreProfile"] + self.base_protocols
         )
+        self.sorted_templates = self._sort_and_filter_models(
+            self.templates, filters=["CoreObjectTemplate"] + self.base_protocols
+        )
 
     def render(self, sync: bool = True) -> str:
         jinja2_env = jinja2.Environment(loader=jinja2.BaseLoader(), trim_blocks=True, lstrip_blocks=True)
-        jinja2_env.filters["inheritance"] = self._jinja2_filter_inheritance
         jinja2_env.filters["render_attribute"] = self._jinja2_filter_render_attribute
         jinja2_env.filters["render_relationship"] = self._jinja2_filter_render_relationship
+        jinja2_env.filters["syncify"] = self._jinja2_filter_syncify
 
         template = jinja2_env.from_string(load_template())
         return template.render(
             generics=self.sorted_generics,
             nodes=self.sorted_nodes,
             profiles=self.sorted_profiles,
+            templates=self.sorted_templates,
             base_protocols=self.base_protocols,
             core_node_name="CoreNodeSync" if sync else "CoreNode",
             sync=sync,
         )
 
     @staticmethod
-    def _jinja2_filter_inheritance(value: dict[str, Any]) -> str:
-        inherit_from: list[str] = value.get("inherit_from", [])
+    def _jinja2_filter_syncify(value: str | list, sync: bool = False) -> str | list:
+        """Filter to help with the convertion to sync
+
+        If a string is provided, append Sync to the end of the string
+        If a list is provided, search for CoreNode and replace it with CoreNodeSync
+        """
+        if not sync:
+            return value
+
+        if isinstance(value, str):
+            return f"{value}Sync"
+
+        if isinstance(value, list):
+            return [f"{item}Sync" if item == "CoreNode" else item for item in value]
 
-        if not inherit_from:
-            return "CoreNode"
-        return ", ".join(inherit_from)
+        return value
 
     @staticmethod
     def _jinja2_filter_render_attribute(value: AttributeSchemaAPI) -> str:

infrahub_sdk/protocols_generator/template.j2

@@ -68,7 +68,7 @@ class {{ generic.namespace + generic.name }}({{core_node_name}}):
 {% for node in nodes %}
 
 
-class {{ node.namespace + node.name }}({{ node.inherit_from | join(", ") or core_node_name }}):
+class {{ node.namespace + node.name }}({{ node.inherit_from | syncify(sync) | join(", ") or core_node_name }}):
     {% if not node.attributes|default([]) and not node.relationships|default([]) %}
     pass
     {% endif %}
@@ -91,7 +91,7 @@ class {{ node.namespace + node.name }}({{ node.inherit_from | join(", ") or core
 {% for node in profiles %}
 
 
-class {{ node.namespace + node.name }}({{ node.inherit_from | join(", ") or core_node_name }}):
+class {{ node.namespace + node.name }}({{ node.inherit_from | syncify(sync) | join(", ") or core_node_name }}):
     {% if not node.attributes|default([]) and not node.relationships|default([]) %}
     pass
     {% endif %}
@@ -102,12 +102,21 @@ class {{ node.namespace + node.name }}({{ node.inherit_from | join(", ") or core
     {{ relationship | render_relationship(sync) }}
     {% endfor %}
     {% if node.hierarchical | default(false) %}
-    {% if sync %}
-    parent: RelatedNodeSync
-    children: RelationshipManagerSync
-    {% else %}
-    parent: RelatedNode
-    children: RelationshipManager
-    {% endif %}
+    parent: {{ "RelatedNode" | sincify(sync)}}
+    children: {{ "RelationshipManager" | sincify(sync)}}
     {% endif %}
 {% endfor %}
+{% for node in templates %}
+
+
+class {{ node.namespace + node.name }}({{ node.inherit_from | syncify(sync) | join(", ") or core_node_name }}):
+    {% if not node.attributes|default([]) and not node.relationships|default([]) %}
+    pass
+    {% endif %}
+    {% for attribute in node.attributes|default([]) %}
+    {{ attribute | render_attribute }}
+    {% endfor %}
+    {% for relationship in node.relationships|default([]) %}
+    {{ relationship | render_relationship(sync) }}
+    {% endfor %}
+{% endfor %}