Merge pull request #242 from mirumee/no_reimports_plugin

No reimports plugin

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## UNRELEASED
 
 - Removed `model_rebuild` calls for generated input, fragment and result models. 
+- Added `NoReimportsPlugin` that makes the `__init__.py` of generated client package empty.
 
 
 ## 0.10.0 (2023-11-15)

README.md

@@ -103,6 +103,8 @@ plugins = ["ariadne_codegen.contrib.extract_operations.ExtractOperationsPlugin"]
 operations_module_name = "custom_operations_module_name"
 ```
 
+- [`ariadne_codegen.contrib.extract_operations.NoReimportsPlugin`](ariadne_codegen/contrib/no_reimports.py) - This plugin removes content of generated `__init__.py`. This is useful in scenarios where generated plugins contain so many Pydantic models that client's eager initialization of entire package on first import is very slow.
+
 
 ## Using generated client
 

ariadne_codegen/contrib/__init__.py

@@ -1,4 +1,5 @@
 from .extract_operations import ExtractOperationsPlugin
+from .no_reimports import NoReimportsPlugin
 from .shorter_results import ShorterResultsPlugin
 
-__all__ = ["ExtractOperationsPlugin", "ShorterResultsPlugin"]
+__all__ = ["ExtractOperationsPlugin", "NoReimportsPlugin", "ShorterResultsPlugin"]

ariadne_codegen/contrib/extract_operations.py

@@ -41,24 +41,25 @@ def __init__(self, schema: GraphQLSchema, config_dict: Dict) -> None:
         self._operations_variables: Dict[str, str] = {}
 
     def generate_init_module(self, module: ast.Module) -> ast.Module:
-        variables_names = list(self._operations_variables.values())
-        module.body.insert(
-            0,
-            generate_import_from(
-                names=variables_names,
-                from_=self.operations_module_name,
-                level=1,
-            ),
-        )
-        all_assign = cast(ast.Assign, module.body[-1])
-        already_imported_names = [
-            cast(ast.Constant, const).value
-            for const in cast(ast.List, all_assign.value).elts
-        ]
-        cast(ast.List, all_assign.value).elts = [
-            generate_constant(name)
-            for name in sorted(already_imported_names + variables_names)
-        ]
+        if module.body:
+            variables_names = list(self._operations_variables.values())
+            module.body.insert(
+                0,
+                generate_import_from(
+                    names=variables_names,
+                    from_=self.operations_module_name,
+                    level=1,
+                ),
+            )
+            all_assign = cast(ast.Assign, module.body[-1])
+            already_imported_names = [
+                cast(ast.Constant, const).value
+                for const in cast(ast.List, all_assign.value).elts
+            ]
+            cast(ast.List, all_assign.value).elts = [
+                generate_constant(name)
+                for name in sorted(already_imported_names + variables_names)
+            ]
 
         self._generate_operations_module()
         return ast.fix_missing_locations(module)

ariadne_codegen/contrib/no_reimports.py

@@ -0,0 +1,9 @@
+import ast
+
+from ariadne_codegen.plugins.base import Plugin
+
+
+class NoReimportsPlugin(Plugin):
+    def generate_init_module(self, module: ast.Module) -> ast.Module:
+        module.body = []
+        return module

tests/contrib/test_extract_operations.py

@@ -87,6 +87,14 @@ def test_generate_init_module_returns_module_with_added_import(
     )
 
 
+def test_generate_init_module_doesnt_add_import_if_module_has_no_body(config_dict):
+    plugin = ExtractOperationsPlugin(schema=None, config_dict=config_dict)
+
+    modified_module = plugin.generate_init_module(ast.Module(body=[], type_ignores=[]))
+
+    assert not modified_module.body
+
+
 def test_generate_init_module_creates_operations_file(config_dict, empty_init_module):
     node = parse("query testXyz { xyz }").definitions[0]
     plugin = ExtractOperationsPlugin(schema=None, config_dict=config_dict)

tests/contrib/test_no_reimports.py

@@ -0,0 +1,32 @@
+import ast
+
+import pytest
+
+from ariadne_codegen.contrib.no_reimports import NoReimportsPlugin
+
+
+@pytest.mark.parametrize(
+    "module",
+    [
+        ast.Module(
+            body=[
+                ast.ImportFrom(
+                    module="input_types", names=[ast.alias(name="InputA")], level=1
+                ),
+                ast.Assign(
+                    targets=[ast.Name(id="__all__")],
+                    value=ast.List(elts=[ast.Constant(value="InputA")]),
+                ),
+            ],
+            type_ignores=[],
+        ),
+        ast.Module(body=[], type_ignores=[]),
+    ],
+)
+def test_generate_init_module_returns_module_without_body(module):
+    plugin = NoReimportsPlugin(schema=None, config_dict={})
+
+    modified_module = plugin.generate_init_module(module)
+
+    assert isinstance(modified_module, ast.Module)
+    assert not modified_module.body