From 9cc9412d1865fc71ec0ec5f9e2885a126436f77b Mon Sep 17 00:00:00 2001 From: Max Chesterfield Date: Wed, 1 Apr 2026 15:02:55 +1100 Subject: [PATCH 1/2] allow plugins to hook into creation of custom operations Signed-off-by: Max Chesterfield --- .../client_generators/custom_operation.py | 6 ++++-- ariadne_codegen/plugins/base.py | 8 ++++++++ ariadne_codegen/plugins/manager.py | 12 ++++++++++++ 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/ariadne_codegen/client_generators/custom_operation.py b/ariadne_codegen/client_generators/custom_operation.py index 805ac32b..03d72abe 100644 --- a/ariadne_codegen/client_generators/custom_operation.py +++ b/ariadne_codegen/client_generators/custom_operation.py @@ -103,7 +103,7 @@ def generate(self) -> ast.Module: + cast(list[ast.stmt], self._type_imports) + [self._class_def], ) - return module + return self.plugin_manager.generate_custom_module(module) def _add_import(self, import_: Optional[ast.ImportFrom] = None): """Adds an import statement to the list avoiding duplicates.""" @@ -155,7 +155,7 @@ def _generate_method( return_arguments_keys, return_arguments_values ) - return generate_method_definition( + method = generate_method_definition( name=process_name( operation_name, convert_to_snake_case=self.convert_to_snake_case, @@ -182,6 +182,8 @@ def _generate_method( ], decorator_list=[generate_name("classmethod")], ) + + return self.plugin_manager.generate_custom_method(method) def _get_return_type_and_from(self, final_type): """ diff --git a/ariadne_codegen/plugins/base.py b/ariadne_codegen/plugins/base.py index ad29b540..7258118e 100644 --- a/ariadne_codegen/plugins/base.py +++ b/ariadne_codegen/plugins/base.py @@ -145,3 +145,11 @@ def get_file_comment( self, comment: str, code: str, source: Optional[str] = None ) -> str: return comment + + def generate_custom_module( + self, module + ) -> ast.Module: + return module + + def generate_custom_method(self, method_def: ast.FunctionDef) -> ast.FunctionDef: + return method_def diff --git a/ariadne_codegen/plugins/manager.py b/ariadne_codegen/plugins/manager.py index 243e60bb..bdd26075 100644 --- a/ariadne_codegen/plugins/manager.py +++ b/ariadne_codegen/plugins/manager.py @@ -213,3 +213,15 @@ def get_file_comment( return self._apply_plugins_on_object( "get_file_comment", comment, code=code, source=source ) + + def generate_custom_module(self, module: ast.Module) -> ast.Module: + return self._apply_plugins_on_object( + "generate_custom_module", + module + ) + + def generate_custom_method(self, method_def: ast.FunctionDef) -> ast.FunctionDef: + return self._apply_plugins_on_object( + "generate_custom_method", + method_def + ) From b321b6052c504a49e3128c1c6df9c390e1ddb97d Mon Sep 17 00:00:00 2001 From: Max Chesterfield Date: Wed, 1 Apr 2026 15:36:03 +1100 Subject: [PATCH 2/2] tests --- .../client_generators/custom_operation.py | 10 ++++-- tests/plugins/test_manager.py | 35 +++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/ariadne_codegen/client_generators/custom_operation.py b/ariadne_codegen/client_generators/custom_operation.py index 03d72abe..da9d8aae 100644 --- a/ariadne_codegen/client_generators/custom_operation.py +++ b/ariadne_codegen/client_generators/custom_operation.py @@ -103,7 +103,9 @@ def generate(self) -> ast.Module: + cast(list[ast.stmt], self._type_imports) + [self._class_def], ) - return self.plugin_manager.generate_custom_module(module) + if self.plugin_manager: + return self.plugin_manager.generate_custom_module(module) + return module def _add_import(self, import_: Optional[ast.ImportFrom] = None): """Adds an import statement to the list avoiding duplicates.""" @@ -182,8 +184,10 @@ def _generate_method( ], decorator_list=[generate_name("classmethod")], ) - - return self.plugin_manager.generate_custom_method(method) + + if self.plugin_manager: + return self.plugin_manager.generate_custom_method(method) + return method def _get_return_type_and_from(self, final_type): """ diff --git a/tests/plugins/test_manager.py b/tests/plugins/test_manager.py index 8ee4d431..ead72f32 100644 --- a/tests/plugins/test_manager.py +++ b/tests/plugins/test_manager.py @@ -411,3 +411,38 @@ def test_get_file_comment_calls_plugins_get_file_comment( plugin1, plugin2 = plugin_manager_with_mocked_plugins.plugins assert plugin1.get_file_comment.called assert plugin2.get_file_comment.called + +def test_generate_custom_module_calls_plugins_generate_custom_module( + plugin_manager_with_mocked_plugins, +): + plugin_manager_with_mocked_plugins.generate_custom_module( + ast.Module(body=[], type_ignores=[]), + ) + + plugin1, plugin2 = plugin_manager_with_mocked_plugins.plugins + assert plugin1.generate_custom_module.called + assert plugin2.generate_custom_module.called + +def test_generate_custom_method_calls_plugins_generate_custom_method( + plugin_manager_with_mocked_plugins, +): + plugin_manager_with_mocked_plugins.generate_custom_method( + ast.FunctionDef( + name="", + args=ast.arguments( + posonlyargs=[], + args=[], + vararg=None, + kwonlyargs=[], + kw_defaults=[], + kwarg=None, + defaults=[], + ), + body=[], + decorator_list=[], + ) + ) + + plugin1, plugin2 = plugin_manager_with_mocked_plugins.plugins + assert plugin1.generate_custom_method.called + assert plugin2.generate_custom_method.called