fix(python): handle non union self referencing schema dependencies (#…

…4604)

generators/python/sdk/versions.yml

@@ -1,4 +1,13 @@
 # For unreleased changes, use unreleased.yml
+- version: 4.0.0-rc5
+  irVersion: 53
+  changelogEntry:
+    - type: fix
+      summary: | 
+        Pydantic models now call update forward refs on non-uion circular references. This 
+        prevents runtime errors in certain cases where types self reference itself through 
+        a union. 
+
 - version: 4.0.0-rc4
   irVersion: 53
   changelogEntry:

generators/python/src/fern_python/generators/context/pydantic_generator_context.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Callable, List, Optional, Set
+from typing import Callable, Dict, List, Optional, Set
 
 import fern.ir.resources as ir_types
 from fern.generator_exec import GeneratorConfig
+from ordered_set import OrderedSet
 
 from fern_python.codegen import AST, Filepath
 from fern_python.declaration_referencer import AbstractDeclarationReferencer
@@ -72,6 +73,10 @@ def does_circularly_reference_itself(self, type_id: ir_types.TypeId) -> bool:
     def get_non_union_circular_references(self) -> Set[ir_types.TypeId]:
         ...
 
+    @abstractmethod
+    def get_non_union_self_referencing_dependencies_from_types(self) -> Dict[ir_types.TypeId, Set[ir_types.TypeId]]:
+        ...
+
     @abstractmethod
     def do_types_reference_each_other(self, a: ir_types.TypeId, b: ir_types.TypeId) -> bool:
         ...
@@ -90,6 +95,10 @@ def does_type_reference_reference_other_type(
     def get_referenced_types(self, type_id: ir_types.TypeId) -> Set[ir_types.TypeId]:
         ...
 
+    @abstractmethod
+    def get_referenced_types_ordered(self, type_id: ir_types.TypeId) -> OrderedSet[ir_types.TypeId]:
+        ...
+
     @abstractmethod
     def get_class_name_for_type_id(self, type_id: ir_types.TypeId, as_request: bool) -> str:
         ...

generators/python/src/fern_python/generators/context/pydantic_generator_context_impl.py

@@ -1,7 +1,8 @@
-from typing import Callable, List, Optional, Set
+from typing import Callable, Dict, List, Optional, Set
 
 import fern.ir.resources as ir_types
 from fern.generator_exec import GeneratorConfig
+from ordered_set import OrderedSet
 
 from fern_python.codegen import AST, Filepath
 from fern_python.declaration_referencer import AbstractDeclarationReferencer
@@ -53,6 +54,20 @@ def __init__(
             ):
                 self._non_union_self_referencing_type_ids.add(id)
 
+        self._types_with_non_union_self_referencing_dependencies: Dict[ir_types.TypeId, Set[ir_types.TypeId]] = dict()
+        for id, type in self.ir.types.items():
+            for referenced_id in type.referenced_types:
+                referenced_type = self.ir.types[referenced_id]
+                if (
+                    referenced_type.shape.get_as_union().type != "union"
+                    and referenced_type.shape.get_as_union().type != "undiscriminatedUnion"
+                ):
+                    # This referenced type is self-referential
+                    if referenced_id in referenced_type.referenced_types:
+                        if self._types_with_non_union_self_referencing_dependencies.get(id) is None:
+                            self._types_with_non_union_self_referencing_dependencies[id] = set()
+                        self._types_with_non_union_self_referencing_dependencies[id].add(referenced_id)
+
     def get_module_path_in_project(self, module_path: AST.ModulePath) -> AST.ModulePath:
         return self._project_module_path + module_path
 
@@ -147,6 +162,10 @@ def does_circularly_reference_itself(self, type_id: ir_types.TypeId) -> bool:
     def get_non_union_circular_references(self) -> Set[ir_types.TypeId]:
         return self._non_union_self_referencing_type_ids
 
+    # This map goes from every non union type to a list of referenced types that circularly reference themselves
+    def get_non_union_self_referencing_dependencies_from_types(self) -> Dict[ir_types.TypeId, Set[ir_types.TypeId]]:
+        return self._types_with_non_union_self_referencing_dependencies
+
     def do_types_reference_each_other(self, a: ir_types.TypeId, b: ir_types.TypeId) -> bool:
         return self.does_type_reference_other_type(a, b) and self.does_type_reference_other_type(b, a)
 
@@ -163,6 +182,10 @@ def get_referenced_types(self, type_id: ir_types.TypeId) -> Set[ir_types.TypeId]
         declaration = self.ir.types[type_id]
         return self.get_referenced_types_of_type_declaration(declaration)
 
+    def get_referenced_types_ordered(self, type_id: ir_types.TypeId) -> OrderedSet[ir_types.TypeId]:
+        declaration = self.ir.types[type_id]
+        return OrderedSet(list(sorted(self.get_referenced_types_of_type_declaration(declaration))))
+
     def get_declaration_for_type_id(
         self,
         type_id: ir_types.TypeId,

generators/python/src/fern_python/generators/pydantic_model/fern_aware_pydantic_model.py

@@ -76,15 +76,6 @@ def __init__(
             if extends is not None
             else []
         )
-        # Acknowledge forward refs for extended models as well
-        for extended_type in extends:
-            type_id_to_reference = self._type_id_for_forward_ref()
-            if type_id_to_reference is not None and context.does_type_reference_other_type(
-                type_id=extended_type.type_id, other_type_id=type_id_to_reference
-            ):
-                # While we don't want to string reference the extended model, we still want to rebuild the model
-                self._model_contains_forward_refs = True
-                break
 
         models_to_extend.extend(extends_crs)
         self._pydantic_model = PydanticModel(
@@ -144,6 +135,18 @@ def add_field(
 
         return field
 
+    def _add_update_forward_ref_for_transitive_circular_dependency(self, type_id: ir_types.TypeId) -> None:
+        # Get self-referencing types, if the type you're viewing is self-referencing (and not a union), then we need to update the forward refs
+        if type_id in self._context.get_non_union_circular_references():
+            # We update the current model's forward refs independently
+            if self._type_name is not None and type_id == self._type_name.type_id:
+                return
+
+            class_reference = self._context.get_class_reference_for_type_id(
+                type_id, as_request=False, must_import_after_current_declaration=lambda _: False
+            )
+            self._pydantic_model.update_forward_refs_for_given_model(class_reference)
+
     def add_private_instance_field_unsafe(
         self, name: str, type_hint: AST.TypeHint, default_factory: AST.Expression
     ) -> None:
@@ -192,10 +195,6 @@ def _must_import_after_current_declaration(self, type_name: ir_types.DeclaredTyp
                 type_id=type_name.type_id, other_type_id=type_id_to_reference
             )
 
-        is_referencing_circular_reference = type_name.type_id in self._context.get_non_union_circular_references()
-        if is_referencing_circular_reference:
-            should_import_after = is_referencing_circular_reference
-
         if should_import_after:
             self._model_contains_forward_refs = True
 
@@ -268,6 +267,21 @@ def finish(self) -> None:
         if self._model_contains_forward_refs or self._force_update_forward_refs:
             self._pydantic_model.update_forward_refs()
 
+        # Acknowledge forward refs for extended models as well
+        for extended_type in self._extends:
+            type_id_to_reference = self._type_id_for_forward_ref()
+            if type_id_to_reference is not None and self._context.does_type_reference_other_type(
+                type_id=extended_type.type_id, other_type_id=type_id_to_reference
+            ):
+                # While we don't want to string reference the extended model, we still want to rebuild the model
+                self._model_contains_forward_refs = True
+                break
+
+        # Now take any transitive circular dependencies and add them as ghost references and update their forward refs
+        if self._type_name is not None:
+            for referenced_type in self._context.get_referenced_types_ordered(self._type_name.type_id):
+                self._add_update_forward_ref_for_transitive_circular_dependency(referenced_type)
+
         self._pydantic_model.finish()
 
     def _get_validators_generator(self) -> ValidatorsGenerator:

generators/python/src/fern_python/pydantic_codegen/pydantic_model.py

@@ -364,6 +364,16 @@ def update_forward_refs(self) -> None:
             )
         )
 
+    def update_forward_refs_for_given_model(self, given_model: AST.ClassReference) -> None:
+        self._source_file.add_footer_expression(
+            AST.Expression(
+                AST.FunctionInvocation(
+                    function_definition=self._update_forward_ref_function_reference,
+                    args=[AST.Expression(given_model)],
+                )
+            )
+        )
+
     def _get_config_class(self) -> Optional[AST.ClassDeclaration]:
         config = AST.ClassDeclaration(name="Config")
 

packages/seed/src/commands/run/runWithCustomFixture.ts

@@ -118,7 +118,7 @@ function getGeneratorGroup({
     for (const group of groups ?? []) {
         for (const generator of group.generators) {
             if (generator.name === image) {
-                const invocation = { ...generator, absolutePathToLocalOutput: absolutePathToOutput };
+                const invocation = { ...generator, absolutePathToLocalOutput: absolutePathToOutput, version: "latest" };
                 return {
                     group: {
                         ...group,