Merge pull request #283 from bombsimon/fix/input-type-model-rebuild

fix: Include `model_rebuild` for input types as well

CHANGELOG.md

@@ -1,5 +1,10 @@
 # CHANGELOG
 
+## 0.14.0 (Unreleased)
+
+- Re-added `model_rebuild` calls for input types with forward references.
+
+
 ## 0.13.0 (2024-03-4)
 
 - Fixed `str_to_snake_case` utility to capture fully capitalized words followed by an underscore.

ariadne_codegen/client_generators/input_types.py

@@ -13,10 +13,13 @@
  generate_ann_assign,
  generate_class_def,
  generate_constant,
+ generate_expr,
  generate_import_from,
  generate_keyword,
+ generate_method_call,
  generate_module,
  generate_pydantic_field,
+ model_has_forward_refs,
 )
 from ..plugins.manager import PluginManager
 from ..utils import process_name
@@ -28,6 +31,7 @@
  BASE_MODEL_IMPORT,
  FIELD_CLASS,
  LIST,
+ MODEL_REBUILD_METHOD,
  OPTIONAL,
  PLAIN_SERIALIZER,
  PYDANTIC_MODULE,
@@ -85,8 +89,16 @@ def generate(self, types_to_include: Optional[List[str]] = None) -> ast.Module:
  scalar_data = self.custom_scalars[scalar_name]
  self._imports.extend(generate_scalar_imports(scalar_data))
 
- module_body = cast(List[ast.stmt], self._imports) + cast(
- List[ast.stmt], class_defs
+ model_rebuild_calls = [
+ generate_expr(generate_method_call(class_def.name, MODEL_REBUILD_METHOD))
+ for class_def in class_defs
+ if model_has_forward_refs(class_def)
+ ]
+
+ module_body = (
+ cast(List[ast.stmt], self._imports)
+ + cast(List[ast.stmt], class_defs)
+ + cast(List[ast.stmt], model_rebuild_calls)
  )
  module = generate_module(body=module_body)
 

ariadne_codegen/client_generators/result_types.py

@@ -40,6 +40,7 @@
  generate_module,
  generate_pass,
  generate_pydantic_field,
+ model_has_forward_refs,
 )
 from ..exceptions import NotSupported, ParsingError
 from ..plugins.manager import PluginManager
@@ -158,7 +159,7 @@ def generate(self) -> ast.Module:
  model_rebuild_calls = [
  generate_expr(generate_method_call(class_def.name, MODEL_REBUILD_METHOD))
  for class_def in self._class_defs
- if self.include_model_rebuild(class_def)
+ if model_has_forward_refs(class_def)
  ]
 
  module_body = (
@@ -174,11 +175,6 @@ def generate(self) -> ast.Module:
  )
  return module
 
- def include_model_rebuild(self, class_def: ast.ClassDef) -> bool:
- visitor = ClassDefNamesVisitor()
- visitor.visit(class_def)
- return visitor.found_name_with_quote
-
  def get_imports(self) -> List[ast.ImportFrom]:
  return self._imports
 
@@ -576,19 +572,3 @@ def enter_field(node: FieldNode, *_args: Any) -> FieldNode:
  copied_node = deepcopy(node)
  visit(copied_node, RemoveMixinVisitor())
  return copied_node
-
-
-class ClassDefNamesVisitor(ast.NodeVisitor):
- def __init__(self):
- self.found_name_with_quote = False
-
- def visit_Name(self, node): # pylint: disable=C0103
- if '"' in node.id:
- self.found_name_with_quote = True
- self.generic_visit(node)
-
- def visit_Subscript(self, node): # pylint: disable=C0103
- if isinstance(node.value, ast.Name) and node.value.id == "Literal":
- return
-
- self.generic_visit(node)

ariadne_codegen/codegen.py

@@ -332,3 +332,25 @@ def generate_yield(value: Optional[ast.expr] = None) -> ast.Yield:
 
 def generate_pass() -> ast.Pass:
  return ast.Pass()
+
+
+def model_has_forward_refs(class_def: ast.ClassDef) -> bool:
+ visitor = ClassDefNamesVisitor()
+ visitor.visit(class_def)
+ return visitor.found_name_with_quote
+
+
+class ClassDefNamesVisitor(ast.NodeVisitor):
+ def __init__(self):
+ self.found_name_with_quote = False
+
+ def visit_Name(self, node): # pylint: disable=C0103
+ if '"' in node.id:
+ self.found_name_with_quote = True
+ self.generic_visit(node)
+
+ def visit_Subscript(self, node): # pylint: disable=C0103
+ if isinstance(node.value, ast.Name) and node.value.id == "Literal":
+ return
+
+ self.generic_visit(node)

tests/main/clients/example/expected_client/input_types.py

@@ -46,3 +46,7 @@ class NotificationsPreferencesInput(BaseModel):
  receive_push_notifications: bool = Field(alias="receivePushNotifications")
  receive_sms: bool = Field(alias="receiveSms")
  title: str
+
+
+UserCreateInput.model_rebuild()
+UserPreferencesInput.model_rebuild()

tests/main/clients/only_used_inputs_and_enums/expected_client/input_types.py

@@ -26,3 +26,8 @@ class InputAB(BaseModel):
 
 class InputE(BaseModel):
  val: EnumE
+
+
+InputA.model_rebuild()
+InputAA.model_rebuild()
+InputAB.model_rebuild()