Implement feedback from PR review

Author: dgarros

docs/docs/python-sdk/guides/python-typing.mdx

@@ -100,7 +100,7 @@ from lib.protocols import MyOwnObject
 my_object = client.get(MyOwnObject, name__value="example")
 ```
 
-> if you don't have your own Python module, it's possible to use relative path by having the `procotols.py` in the same directory as your script/transform/generator
+> if you don't have your own Python module, it's possible to use relative path by having the `protocols.py` in the same directory as your script/transform/generator
 
 ## Generating Pydantic models from GraphQL queries
 
@@ -110,7 +110,7 @@ When working with GraphQL queries, you can generate type-safe Pydantic models th
 
 Generated Pydantic models from GraphQL queries offer several important benefits:
 
-- **Type Safety**: Catch type errors at development time instead of runtime
+- **Type Safety**: Catch type errors during development time instead of at runtime
 - **IDE Support**: Get autocomplete, type hints, and better IntelliSense in your IDE  
 - **Documentation**: Generated models serve as living documentation of your GraphQL API
 - **Validation**: Automatic validation of query responses against the expected schema

infrahub_sdk/ctl/graphql.py

@@ -47,7 +47,7 @@ def find_gql_files(query_path: Path) -> list[Path]:
         List of Path objects for all .gql files found
     """
     if not query_path.exists():
-        raise FileNotFoundError(f"Directory not found: {query_path}")
+        raise FileNotFoundError(f"File or directory not found: {query_path}")
 
     if not query_path.is_dir() and query_path.is_file():
         return [query_path]
@@ -105,15 +105,19 @@ async def export_schema(
     """Export the GraphQL schema to a file."""
 
     client = initialize_client()
-    response = await client._get(url=f"{client.address}/schema.graphql")
-    destination.write_text(response.text)
+    schema_text = await client.schema.get_graphql_schema()
+
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    destination.write_text(schema_text)
     console.print(f"[green]Schema exported to {destination}")
 
 
 @app.command()
 @catch_exception(console=console)
 async def generate_return_types(
-    query: Optional[Path] = typer.Argument(None, help="Location of the GraphQL query file(s)."),
+    query: Optional[Path] = typer.Argument(
+        None, help="Location of the GraphQL query file(s). Defaults to current directory if not specified."
+    ),
     schema: Path = typer.Option("schema.graphql", help="Path to the GraphQL schema file."),
     _: str = CONFIG_PARAM,
 ) -> None:
@@ -144,7 +148,7 @@ async def generate_return_types(
             try:
                 definitions = get_graphql_query(queries_path=gql_file, schema=graphql_schema)
             except ValueError as exc:
-                print(f"Error generating result types for {gql_file}: {exc}")
+                console.print(f"[red]Error generating result types for {gql_file}: {exc}")
                 continue
             queries = filter_operations_definitions(definitions)
             fragments = filter_fragments_definitions(definitions)
@@ -161,11 +165,17 @@ async def generate_return_types(
                 plugin_manager=plugin_manager,
             )
 
+            parsing_failed = False
             try:
                 for query_operation in queries:
                     package_generator.add_operation(query_operation)
             except ParsingError as exc:
                 console.print(f"[red]Unable to process {gql_file.name}: {exc}")
+                parsing_failed = True
+
+            if parsing_failed:
+                continue
+
             module_fragment = package_generator.fragments_generator.generate()
 
             generate_result_types(directory=directory, package=package_generator, fragment=module_fragment)

infrahub_sdk/graphql/plugin.py

@@ -14,8 +14,9 @@ class FutureAnnotationPlugin(Plugin):
     def insert_future_annotation(module: ast.Module) -> ast.Module:
         # First check if the future annotation is already present
         for item in module.body:
-            if isinstance(item, ast.ImportFrom) and item.module == "__future__" and item.names[0].name == "annotations":
-                return module
+            if isinstance(item, ast.ImportFrom) and item.module == "__future__":
+                if any(alias.name == "annotations" for alias in item.names):
+                    return module
 
         module.body.insert(0, ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0))
         return module
@@ -68,13 +69,11 @@ def find_base_model_index(module: ast.Module) -> int:
         for idx, item in enumerate(module.body):
             if isinstance(item, ast.ImportFrom) and item.module == "base_model":
                 return idx
-        return -1
+        raise ValueError("BaseModel not found in module")
 
     @classmethod
     def replace_base_model_import(cls, module: ast.Module) -> ast.Module:
         base_model_index = cls.find_base_model_index(module)
-        if base_model_index == -1:
-            raise ValueError("BaseModel not found in module")
         module.body[base_model_index] = ast.ImportFrom(module="pydantic", names=[ast.alias(name="BaseModel")], level=0)
         return module
 

infrahub_sdk/graphql/utils.py

@@ -11,12 +11,22 @@ def get_class_def_index(module: ast.Module) -> int:
 
 
 def insert_fragments_inline(module: ast.Module, fragment: ast.Module) -> ast.Module:
-    """Insert the Pydantic classes for the fragments inline into the module."""
+    """Insert the Pydantic classes for the fragments inline into the module.
+
+    If no class definitions exist in module, fragments are appended to the end.
+    """
     module_class_def_index = get_class_def_index(module)
 
     fragment_classes: list[ast.ClassDef] = [item for item in fragment.body if isinstance(item, ast.ClassDef)]
-    for idx, item in enumerate(fragment_classes):
-        module.body.insert(module_class_def_index + idx, item)
+
+    # Handle edge case when no class definitions exist
+    if module_class_def_index == -1:
+        # Append fragments to the end of the module
+        module.body.extend(fragment_classes)
+    else:
+        # Insert fragments before the first class definition
+        for idx, item in enumerate(fragment_classes):
+            module.body.insert(module_class_def_index + idx, item)
 
     return module
 

infrahub_sdk/schema/__init__.py

@@ -474,6 +474,25 @@ async def fetch(
 
         return branch_schema.nodes
 
+    async def get_graphql_schema(self, branch: str | None = None) -> str:
+        """Get the GraphQL schema as a string.
+
+        Args:
+            branch: The branch to get the schema for. Defaults to default_branch.
+
+        Returns:
+            The GraphQL schema as a string.
+        """
+        branch = branch or self.client.default_branch
+        url = f"{self.client.address}/schema.graphql"
+
+        response = await self.client._get(url=url)
+
+        if response.status_code != 200:
+            raise ValueError(f"Failed to fetch GraphQL schema: HTTP {response.status_code} - {response.text}")
+
+        return response.text
+
     async def _fetch(self, branch: str, namespaces: list[str] | None = None) -> BranchSchema:
         url_parts = [("branch", branch)]
         if namespaces:
@@ -697,6 +716,25 @@ def fetch(
 
         return branch_schema.nodes
 
+    def get_graphql_schema(self, branch: str | None = None) -> str:
+        """Get the GraphQL schema as a string.
+
+        Args:
+            branch: The branch to get the schema for. Defaults to default_branch.
+
+        Returns:
+            The GraphQL schema as a string.
+        """
+        branch = branch or self.client.default_branch
+        url = f"{self.client.address}/schema.graphql"
+
+        response = self.client._get(url=url)
+
+        if response.status_code != 200:
+            raise ValueError(f"Failed to fetch GraphQL schema: HTTP {response.status_code} - {response.text}")
+
+        return response.text
+
     def _fetch(self, branch: str, namespaces: list[str] | None = None) -> BranchSchema:
         url_parts = [("branch", branch)]
         if namespaces:

tests/unit/sdk/graphql/test_renderer.py

@@ -1,106 +1,6 @@
-from enum import Enum
-
-import pytest
-
 from infrahub_sdk.graphql.renderers import render_input_block, render_query_block
 
 
-class MyStrEnum(str, Enum):
-    VALUE1 = "value1"
-    VALUE2 = "value2"
-
-
-class MyIntEnum(int, Enum):
-    VALUE1 = 12
-    VALUE2 = 24
-
-
-@pytest.fixture
-def query_data_no_filter():
-    data = {
-        "device": {
-            "name": {"value": None},
-            "description": {"value": None},
-            "interfaces": {"name": {"value": None}},
-        }
-    }
-
-    return data
-
-
-@pytest.fixture
-def query_data_alias():
-    data = {
-        "device": {
-            "name": {"@alias": "new_name", "value": None},
-            "description": {"value": {"@alias": "myvalue"}},
-            "interfaces": {"@alias": "myinterfaces", "name": {"value": None}},
-        }
-    }
-
-    return data
-
-
-@pytest.fixture
-def query_data_fragment():
-    data = {
-        "device": {
-            "name": {"value": None},
-            "...on Builtin": {
-                "description": {"value": None},
-                "interfaces": {"name": {"value": None}},
-            },
-        }
-    }
-
-    return data
-
-
-@pytest.fixture
-def query_data_empty_filter():
-    data = {
-        "device": {
-            "@filters": {},
-            "name": {"value": None},
-            "description": {"value": None},
-            "interfaces": {"name": {"value": None}},
-        }
-    }
-
-    return data
-
-
-@pytest.fixture
-def query_data_filters_01():
-    data = {
-        "device": {
-            "@filters": {"name__value": "$name"},
-            "name": {"value": None},
-            "description": {"value": None},
-            "interfaces": {
-                "@filters": {"enabled__value": "$enabled"},
-                "name": {"value": None},
-            },
-        }
-    }
-    return data
-
-
-@pytest.fixture
-def query_data_filters_02():
-    data = {
-        "device": {
-            "@filters": {"name__value": "myname", "integer__value": 44, "enumstr__value": MyStrEnum.VALUE2},
-            "name": {"value": None},
-            "interfaces": {
-                "@filters": {"enabled__value": True, "enumint__value": MyIntEnum.VALUE1},
-                "name": {"value": None},
-            },
-        }
-    }
-    return data
-
-
 def test_render_query_block(query_data_no_filter) -> None:
     lines = render_query_block(data=query_data_no_filter)