diff --git a/ariadne_codegen/client_generators/client.py b/ariadne_codegen/client_generators/client.py index 936f93cc..a0b42904 100644 --- a/ariadne_codegen/client_generators/client.py +++ b/ariadne_codegen/client_generators/client.py @@ -1,4 +1,5 @@ import ast +from copy import deepcopy from typing import Optional, Union, cast from graphql import OperationDefinitionNode, OperationType @@ -819,6 +820,10 @@ def get_variable_names(self, arguments: ast.arguments) -> dict[str, str]: def _add_import(self, import_: Optional[ast.ImportFrom] = None): if not import_: return + # Shared module-level nodes (UNSET_IMPORT, UPLOAD_IMPORT, ...) must never + # reach a generated module: plugins rewrite imports in place, which would + # otherwise corrupt the constant for every later run in this process. + import_ = deepcopy(import_) if self.plugin_manager: import_ = self.plugin_manager.generate_client_import(import_) if import_.names and import_.module: diff --git a/ariadne_codegen/client_generators/fragments.py b/ariadne_codegen/client_generators/fragments.py index 84ae0218..d4a19176 100644 --- a/ariadne_codegen/client_generators/fragments.py +++ b/ariadne_codegen/client_generators/fragments.py @@ -22,6 +22,7 @@ def __init__( plugin_manager: Optional[PluginManager] = None, default_optional_fields_to_none: bool = False, include_typename: bool = True, + defer_model_build: bool = False, ) -> None: self.schema = schema self.enums_module_name = enums_module_name @@ -32,6 +33,7 @@ def __init__( self.plugin_manager = plugin_manager self.default_optional_fields_to_none = default_optional_fields_to_none self.include_typename = include_typename + self.defer_model_build = defer_model_build self._fragments_names = set(self.fragments_definitions.keys()) self._generated_public_names: list[str] = [] @@ -130,6 +132,8 @@ def visit(name): def _get_model_rebuild_calls( self, top_level_fragments_names: list[str], class_defs: list[ast.ClassDef] ) -> list[ast.Call]: + if self.defer_model_build: + return [] class_names = [c.name for c in class_defs] sorted_fragments_names = sorted( top_level_fragments_names, key=class_names.index diff --git a/ariadne_codegen/client_generators/input_types.py b/ariadne_codegen/client_generators/input_types.py index 6224f62b..30c6cd2f 100644 --- a/ariadne_codegen/client_generators/input_types.py +++ b/ariadne_codegen/client_generators/input_types.py @@ -1,5 +1,6 @@ import ast from collections import defaultdict +from copy import deepcopy from typing import Optional, cast from graphql import ( @@ -54,18 +55,20 @@ def __init__( convert_to_snake_case: bool = True, custom_scalars: Optional[dict[str, ScalarData]] = None, plugin_manager: Optional[PluginManager] = None, + defer_model_build: bool = False, ) -> None: self.schema = schema self.convert_to_snake_case = convert_to_snake_case self.enums_module = enums_module self.custom_scalars = custom_scalars if custom_scalars else {} self.plugin_manager = plugin_manager + self.defer_model_build = defer_model_build self._imports = [ generate_import_from([OPTIONAL, ANY, UNION, ANNOTATED], TYPING_MODULE), generate_import_from([FIELD_CLASS, PLAIN_SERIALIZER], PYDANTIC_MODULE), - base_model_import, - *([] if upload_import is None else [upload_import]), + deepcopy(base_model_import), + *([] if upload_import is None else [deepcopy(upload_import)]), ] self._dependencies: dict[str, list[str]] = defaultdict(list) self._used_enums: dict[str, list[str]] = defaultdict(list) @@ -88,11 +91,17 @@ 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)) - 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) - ] + model_rebuild_calls = ( + [] + if self.defer_model_build + else [ + 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) diff --git a/ariadne_codegen/client_generators/package.py b/ariadne_codegen/client_generators/package.py index adbd75b3..24ff9a6c 100644 --- a/ariadne_codegen/client_generators/package.py +++ b/ariadne_codegen/client_generators/package.py @@ -1,4 +1,6 @@ import ast +from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path from typing import Optional @@ -14,8 +16,11 @@ from ..plugins.manager import PluginManager from ..settings import ClientSettings, CommentsStrategy from ..utils import ( + _format_code, + add_defer_build_to_base_model, add_extra_to_base_model, - ast_to_str, + ast_to_raw_str, + format_many, process_name, str_to_pascal_case, ) @@ -54,6 +59,17 @@ from .scalars import ScalarData +@dataclass +class _PendingFile: + """A generated module awaiting formatting and a write to disk.""" + + path: Path + code: str + remove_unused_imports: bool + comment_source: Optional[str] = None + plugin_hook: Optional[Callable[[str], str]] = None + + class PackageGenerator: def __init__( self, @@ -99,6 +115,7 @@ def __init__( default_optional_fields_to_none: bool = False, include_typename: bool = True, ignore_extra_fields: bool = True, + defer_model_build: bool = False, ) -> None: self.package_path = Path(target_path) / package_name @@ -155,9 +172,11 @@ def __init__( self.default_optional_fields_to_none = default_optional_fields_to_none self.include_typename = include_typename self.ignore_extra_fields = ignore_extra_fields + self.defer_model_build = defer_model_build self._result_types_files: dict[str, ast.Module] = {} self._generated_files: list[str] = [] + self._pending_files: list[_PendingFile] = [] self._unpacked_fragments: set[str] = set() self._used_enums: list[str] = [] @@ -193,9 +212,56 @@ def generate(self) -> list[str]: self._generate_client() self._generate_enums() self._generate_init() + self._write_pending_files() return sorted(self._generated_files) + def _queue_module( + self, + path: Path, + module: ast.Module, + remove_unused_imports: bool = True, + multiline_strings: bool = False, + comment_source: Optional[str] = None, + plugin_hook: Optional[Callable[[str], str]] = None, + ) -> None: + """Unparse a module now, format and write it once generation is done.""" + self._pending_files.append( + _PendingFile( + path=path, + code=ast_to_raw_str(module, multiline_strings=multiline_strings), + remove_unused_imports=remove_unused_imports, + comment_source=comment_source, + plugin_hook=plugin_hook if self.plugin_manager else None, + ) + ) + self._generated_files.append(path.name) + + def _write_pending_files(self) -> None: + """Format every queued module, then apply comments, plugins and write. + + Modules are batched by their ruff rule selection so that ruff is spawned + a fixed number of times rather than twice per generated file. Plugins + still see formatted code, as they did when each file formatted alone. + """ + for remove_unused_imports in (True, False): + group = [ + pending_file + for pending_file in self._pending_files + if pending_file.remove_unused_imports is remove_unused_imports + ] + formatted = format_many( + [pending_file.code for pending_file in group], + remove_unused_imports=remove_unused_imports, + ) + for pending_file, code in zip(group, formatted, strict=True): + code = self._add_comments_to_code(code, pending_file.comment_source) + if pending_file.plugin_hook: + code = pending_file.plugin_hook(code) + pending_file.path.write_text(code) + + self._pending_files.clear() + def add_operation(self, definition: OperationDefinitionNode): name = definition.name if not name: @@ -223,6 +289,7 @@ def add_operation(self, definition: OperationDefinitionNode): plugin_manager=self.plugin_manager, default_optional_fields_to_none=self.default_optional_fields_to_none, include_typename=self.include_typename, + defer_model_build=self.defer_model_build, ) self._unpacked_fragments = self._unpacked_fragments.union( query_types_generator.get_unpacked_fragments() @@ -283,13 +350,17 @@ def _validate_unique_file_names(self): def _generate_client(self): client_file_path = self.package_path / f"{self.client_file_name}.py" client_module = self.client_generator.generate() - code = self._add_comments_to_code( - ast_to_str(client_module, multiline_strings=True), self.queries_source + self._queue_module( + client_file_path, + client_module, + multiline_strings=True, + comment_source=self.queries_source, + plugin_hook=( + self.plugin_manager.generate_client_code + if self.plugin_manager + else None + ), ) - if self.plugin_manager: - code = self.plugin_manager.generate_client_code(code) - client_file_path.write_text(code) - self._generated_files.append(client_file_path.name) self._used_enums.extend( self.client_generator.arguments_generator.get_used_enums() ) @@ -314,12 +385,15 @@ def _generate_enums(self): else: module = self.enums_generator.generate(types_to_include=self._used_enums) - code = self._add_comments_to_code(ast_to_str(module), self.schema_source) - if self.plugin_manager: - code = self.plugin_manager.generate_enums_code(code) enums_file_path = self.package_path / f"{self.enums_module_name}.py" - enums_file_path.write_text(code) - self._generated_files.append(enums_file_path.name) + self._queue_module( + enums_file_path, + module, + comment_source=self.schema_source, + plugin_hook=( + self.plugin_manager.generate_enums_code if self.plugin_manager else None + ), + ) self.init_generator.add_import( self.enums_generator.get_generated_public_names(), self.enums_module_name, 1 ) @@ -332,11 +406,16 @@ def _generate_input_types(self): module = self.input_types_generator.generate(types_to_include=used_inputs) input_types_file_path = self.package_path / f"{self.input_types_module_name}.py" - code = self._add_comments_to_code(ast_to_str(module), self.schema_source) - if self.plugin_manager: - code = self.plugin_manager.generate_inputs_code(code) - input_types_file_path.write_text(code) - self._generated_files.append(input_types_file_path.name) + self._queue_module( + input_types_file_path, + module, + comment_source=self.schema_source, + plugin_hook=( + self.plugin_manager.generate_inputs_code + if self.plugin_manager + else None + ), + ) self._used_enums.extend(self.input_types_generator.get_used_enums()) self.init_generator.add_import( self.input_types_generator.get_generated_public_names(), @@ -346,12 +425,16 @@ def _generate_input_types(self): def _generate_result_types(self): for file_name, module in self._result_types_files.items(): - file_path = self.package_path / file_name - code = self._add_comments_to_code(ast_to_str(module), self.queries_source) - if self.plugin_manager: - code = self.plugin_manager.generate_result_types_code(code) - file_path.write_text(code) - self._generated_files.append(file_path.name) + self._queue_module( + self.package_path / file_name, + module, + comment_source=self.queries_source, + plugin_hook=( + self.plugin_manager.generate_result_types_code + if self.plugin_manager + else None + ), + ) def _generate_fragments(self): if not set(self.fragments_definitions.keys()).difference( @@ -363,9 +446,7 @@ def _generate_fragments(self): exclude_names=self._unpacked_fragments ) file_path = self.package_path / f"{self.fragments_module_name}.py" - code = self._add_comments_to_code(ast_to_str(module), self.queries_source) - file_path.write_text(code) - self._generated_files.append(file_path.name) + self._queue_module(file_path, module, comment_source=self.queries_source) self._used_enums.extend(self.fragments_generator.get_used_enums()) self.init_generator.add_import( self.fragments_generator.get_generated_public_names(), @@ -381,9 +462,8 @@ def _copy_files(self): } for source_path, target_name in files_to_copy.items(): code = self._add_comments_to_code(source_path.read_text(encoding="utf-8")) - is_base_model = source_path == self.base_model_file_path - if is_base_model and not self.ignore_extra_fields: - code = add_extra_to_base_model(code) + if source_path == self.base_model_file_path: + code = self._rewrite_base_model(code) if self.plugin_manager: code = self.plugin_manager.copy_code(code) target_path = self.package_path / target_name @@ -404,46 +484,62 @@ def _copy_files(self): level=1, ) + def _rewrite_base_model(self, code: str) -> str: + rewritten = code + if not self.ignore_extra_fields: + rewritten = add_extra_to_base_model(rewritten) + if self.defer_model_build: + rewritten = add_defer_build_to_base_model(rewritten) + + if rewritten == code: + return code + # Each rewrite round-trips through `ast.unparse`, which drops the blank + # lines and import order of the copied dependency. + return _format_code(rewritten, remove_unused_imports=False) + def _generate_init(self): init_file_path = self.package_path / "__init__.py" init_module = self.init_generator.generate() - code = self._add_comments_to_code(ast_to_str(init_module, False)) - if self.plugin_manager: - code = self.plugin_manager.generate_init_code(code) - init_file_path.write_text(code) - self._generated_files.append(init_file_path.name) + self._queue_module( + init_file_path, + init_module, + remove_unused_imports=False, + plugin_hook=( + self.plugin_manager.generate_init_code if self.plugin_manager else None + ), + ) def _generate_custom_queries(self): assert self.custom_query_generator is not None - file_path = self.package_path / "custom_queries.py" - module = self.custom_query_generator.generate() - code = self._add_comments_to_code(ast_to_str(module, False)) - file_path.write_text(code) - self._generated_files.append(file_path.name) + self._queue_module( + self.package_path / "custom_queries.py", + self.custom_query_generator.generate(), + remove_unused_imports=False, + ) def _generate_custom_mutations(self): assert self.custom_mutation_generator is not None - file_path = self.package_path / "custom_mutations.py" - module = self.custom_mutation_generator.generate() - code = self._add_comments_to_code(ast_to_str(module, False)) - file_path.write_text(code) - self._generated_files.append(file_path.name) + self._queue_module( + self.package_path / "custom_mutations.py", + self.custom_mutation_generator.generate(), + remove_unused_imports=False, + ) def _generate_custom_fields_typing(self): assert self.custom_fields_typing_generator is not None - file_path = self.package_path / "custom_typing_fields.py" - module = self.custom_fields_typing_generator.generate() - code = self._add_comments_to_code(ast_to_str(module, False)) - file_path.write_text(code) - self._generated_files.append(file_path.name) + self._queue_module( + self.package_path / "custom_typing_fields.py", + self.custom_fields_typing_generator.generate(), + remove_unused_imports=False, + ) def _generate_custom_fields(self): assert self.custom_fields_generator is not None - file_path = self.package_path / "custom_fields.py" - module = self.custom_fields_generator.generate() - code = self._add_comments_to_code(ast_to_str(module, False)) - file_path.write_text(code) - self._generated_files.append(file_path.name) + self._queue_module( + self.package_path / "custom_fields.py", + self.custom_fields_generator.generate(), + remove_unused_imports=False, + ) def get_package_generator( @@ -493,6 +589,7 @@ def get_package_generator( convert_to_snake_case=settings.convert_to_snake_case, custom_scalars=settings.scalars, plugin_manager=plugin_manager, + defer_model_build=settings.defer_model_build, ) fragments_definitions = {f.name.value: f for f in fragments or []} fragments_generator = FragmentsGenerator( @@ -505,6 +602,7 @@ def get_package_generator( plugin_manager=plugin_manager, default_optional_fields_to_none=settings.default_optional_fields_to_none, include_typename=settings.include_typename, + defer_model_build=settings.defer_model_build, ) custom_fields_generator = CustomFieldsGenerator( schema=schema, @@ -589,4 +687,5 @@ def get_package_generator( default_optional_fields_to_none=settings.default_optional_fields_to_none, include_typename=settings.include_typename, ignore_extra_fields=settings.ignore_extra_fields, + defer_model_build=settings.defer_model_build, ) diff --git a/ariadne_codegen/client_generators/result_types.py b/ariadne_codegen/client_generators/result_types.py index 907b05a1..6dd97077 100644 --- a/ariadne_codegen/client_generators/result_types.py +++ b/ariadne_codegen/client_generators/result_types.py @@ -87,6 +87,7 @@ def __init__( plugin_manager: Optional[PluginManager] = None, default_optional_fields_to_none: bool = False, include_typename: bool = True, + defer_model_build: bool = False, ) -> None: self.schema = schema self.operation_definition = operation_definition @@ -103,13 +104,14 @@ def __init__( self.plugin_manager = plugin_manager self.default_optional_fields_to_none = default_optional_fields_to_none self.include_typename = include_typename + self.defer_model_build = defer_model_build self._imports: list[ast.ImportFrom] = [ generate_import_from( [OPTIONAL, UNION, ANY, LITERAL, ANNOTATED], TYPING_MODULE ), generate_import_from([FIELD_CLASS, BEFORE_VALIDATOR], PYDANTIC_MODULE), - base_model_import + deepcopy(base_model_import) or generate_import_from([BASE_MODEL_CLASS_NAME], PYDANTIC_MODULE), ] self._public_names: list[str] = [] @@ -162,11 +164,17 @@ def _get_operation_type_name(self, definition: ExecutableDefinitionNode) -> str: raise NotSupported(f"Not supported operation type: {definition}") 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 model_has_forward_refs(class_def) - ] + model_rebuild_calls = ( + [] + if self.defer_model_build + else [ + generate_expr( + generate_method_call(class_def.name, MODEL_REBUILD_METHOD) + ) + for class_def in self._class_defs + if model_has_forward_refs(class_def) + ] + ) module_body = ( cast(list[ast.stmt], self._imports) diff --git a/ariadne_codegen/settings.py b/ariadne_codegen/settings.py index 2d15a70d..82b71fe4 100644 --- a/ariadne_codegen/settings.py +++ b/ariadne_codegen/settings.py @@ -185,6 +185,7 @@ class ClientSettings(BaseSettings): default_optional_fields_to_none: bool = False include_typename: bool = True ignore_extra_fields: bool = True + defer_model_build: bool = False def __post_init__(self): if not self.queries_path and not self.enable_custom_operations: @@ -305,6 +306,12 @@ def used_settings_message(self) -> str: if self.include_typename else "Not including __typename fields in generated queries." ) + defer_model_build_msg = ( + "Deferring Pydantic model builds to first use " + "(faster import of generated package)." + if self.defer_model_build + else "Building Pydantic models eagerly at import time." + ) introspection_msg = ( self._introspection_settings_message() if self.using_remote_schema else "" ) @@ -326,6 +333,7 @@ def used_settings_message(self) -> str: {snake_case_msg} {async_client_msg} {include_typename_msg} + {defer_model_build_msg} {files_to_include_msg} {plugins_msg} """ diff --git a/ariadne_codegen/utils.py b/ariadne_codegen/utils.py index b52a4d9b..a4b48508 100644 --- a/ariadne_codegen/utils.py +++ b/ariadne_codegen/utils.py @@ -1,9 +1,11 @@ import ast import builtins import re +import shutil import subprocess import sys import tempfile +from functools import lru_cache from keyword import iskeyword from pathlib import Path from textwrap import indent @@ -18,61 +20,59 @@ _TARGET_VERSION = "py310" _SUBPROCESS_TIMEOUT = 30 +# `ruff check --fix` exits 1 when violations it cannot fix remain. That is not an +# error for us: the fixable ones were still applied. +_RUFF_CHECK_OK_RETURN_CODES = (0, 1) -def _format_code(code: str, *, remove_unused_imports: bool = True) -> str: - """Format generated code with ruff: sort imports, remove unused imports, and format. - Uses ``--isolated`` so the output is deterministic regardless of the - user's ruff configuration. +@lru_cache(maxsize=1) +def _ruff_command() -> list[str]: + """Command prefix used to invoke ruff. + + Running the binary directly skips a full Python interpreter startup on every + call, which otherwise dominates generation time for packages with many + modules. Falls back to `python -m ruff` if the binary cannot be located. """ - select_rules = ["I"] - if remove_unused_imports: - select_rules.append("F401") - - with tempfile.NamedTemporaryFile( - mode="w", - suffix=".py", - delete=False, - encoding="utf-8", - ) as f: - f.write(code) - tmp_path = f.name try: - subprocess.run( - [ - sys.executable, - "-m", - "ruff", - "check", - "--fix", - "--isolated", - "--select", - ",".join(select_rules), - "--target-version", - _TARGET_VERSION, - "--line-length", - str(_LINE_LENGTH), - tmp_path, - ], - check=False, - capture_output=True, - timeout=_SUBPROCESS_TIMEOUT, - ) - code = Path(tmp_path).read_text(encoding="utf-8") - finally: - Path(tmp_path).unlink(missing_ok=True) + from ruff.__main__ import find_ruff_bin - result = subprocess.run( - [ - sys.executable, - "-m", - "ruff", - "format", - "--isolated", - "--target-version", - _TARGET_VERSION, - "--line-length", - str(_LINE_LENGTH), + return [find_ruff_bin()] + except (ImportError, FileNotFoundError): + pass + + ruff_path = shutil.which("ruff") + if ruff_path: + return [ruff_path] + + return [sys.executable, "-m", "ruff"] + + +def _ruff_style_args() -> list[str]: + """`--isolated` keeps output deterministic regardless of the user's config.""" + return [ + "--isolated", + "--target-version", + _TARGET_VERSION, + "--line-length", + str(_LINE_LENGTH), + ] + + +def _ruff_select(remove_unused_imports: bool) -> str: + return "I,F401" if remove_unused_imports else "I" + + +def _format_code(code: str, *, remove_unused_imports: bool = True) -> str: + """Format a single module with ruff: sort imports, drop unused ones, format.""" + check = subprocess.run( + _ruff_command() + + ["check", "--fix"] + + _ruff_style_args() + + [ + "--select", + _ruff_select(remove_unused_imports), + "--stdin-filename", + "generated.py", "-", ], input=code, @@ -82,6 +82,20 @@ def _format_code(code: str, *, remove_unused_imports: bool = True) -> str: check=False, timeout=_SUBPROCESS_TIMEOUT, ) + if check.returncode in _RUFF_CHECK_OK_RETURN_CODES: + # An empty result is legitimate (e.g. a module of only unused imports), + # so this must not fall back to `code` on a falsy stdout. + code = check.stdout + + result = subprocess.run( + _ruff_command() + ["format"] + _ruff_style_args() + ["-"], + input=code, + capture_output=True, + text=True, + encoding="utf-8", + check=False, + timeout=_SUBPROCESS_TIMEOUT, + ) if result.returncode != 0: raise RuntimeError( f"ruff format failed (exit code {result.returncode}): {result.stderr}" @@ -89,6 +103,48 @@ def _format_code(code: str, *, remove_unused_imports: bool = True) -> str: return result.stdout +def format_many(codes: list[str], *, remove_unused_imports: bool = True) -> list[str]: + """Format many modules with a single pair of ruff invocations. + + Equivalent to calling `_format_code` on each module, but ruff walks a scratch + directory instead of being spawned twice per module. + """ + if not codes: + return [] + + with tempfile.TemporaryDirectory() as tmp_dir: + paths = [] + for index, code in enumerate(codes): + path = Path(tmp_dir) / f"module_{index}.py" + path.write_text(code, encoding="utf-8") + paths.append(path) + + subprocess.run( + _ruff_command() + + ["check", "--fix", "--no-cache"] + + _ruff_style_args() + + ["--select", _ruff_select(remove_unused_imports), tmp_dir], + capture_output=True, + check=False, + timeout=_SUBPROCESS_TIMEOUT, + ) + + result = subprocess.run( + _ruff_command() + ["format", "--no-cache"] + _ruff_style_args() + [tmp_dir], + capture_output=True, + text=True, + encoding="utf-8", + check=False, + timeout=_SUBPROCESS_TIMEOUT, + ) + if result.returncode != 0: + raise RuntimeError( + f"ruff format failed (exit code {result.returncode}): {result.stderr}" + ) + + return [path.read_text(encoding="utf-8") for path in paths] + + PYDANTIC_RESERVED_FIELD_NAMES = [ name for name in dir(BaseModel) if not name.startswith("_") ] @@ -102,18 +158,30 @@ def _is_builtin_type_name(name: str) -> bool: return isinstance(value, type) -def ast_to_str( +def ast_to_raw_str( ast_obj: ast.AST, - remove_unused_imports: bool = True, multiline_strings: bool = False, multiline_strings_offset: int = 4, ) -> str: - """Convert ast object into string.""" + """Convert ast object into string, without running ruff over it.""" code = ast.unparse(ast_obj) code = remove_blank_line_between_class_and_content(code) if multiline_strings: code = format_multiline_strings(code, offset=multiline_strings_offset) - return _format_code(code, remove_unused_imports=remove_unused_imports) + return code + + +def ast_to_str( + ast_obj: ast.AST, + remove_unused_imports: bool = True, + multiline_strings: bool = False, + multiline_strings_offset: int = 4, +) -> str: + """Convert ast object into string.""" + return _format_code( + ast_to_raw_str(ast_obj, multiline_strings, multiline_strings_offset), + remove_unused_imports=remove_unused_imports, + ) def remove_blank_line_between_class_and_content(code: str) -> str: @@ -223,8 +291,27 @@ def process_name( return processed_name -def add_extra_to_base_model(code: str) -> str: - "Adds `extra='forbid'` to the ConfigDict in BaseModel if not already present." +def _split_leading_comments(code: str) -> tuple[str, str]: + """Separate a module's leading comment block from its code. + + `ast.unparse` drops comments, which would otherwise discard the header that + `_add_comments_to_code` puts on every generated file. + """ + lines = code.splitlines(keepends=True) + for index, line in enumerate(lines): + stripped = line.strip() + if stripped and not stripped.startswith("#"): + return "".join(lines[:index]), "".join(lines[index:]) + return "", code + + +def _set_base_model_config_kwarg(code: str, arg: str, value: ast.expr) -> str: + """Set a keyword argument on the ``ConfigDict`` call in ``BaseModel``. + + Does nothing if the keyword is already present, so explicit user values are + never overwritten. + """ + header, code = _split_leading_comments(code) tree = ast.parse(code) for node in tree.body: if not isinstance(node, ast.ClassDef): @@ -241,9 +328,17 @@ def add_extra_to_base_model(code: str) -> str: continue if call.func.id != "ConfigDict": continue - if not any(kw.arg == "extra" for kw in call.keywords): - call.keywords.append( - ast.keyword(arg="extra", value=ast.Constant("forbid")) - ) + if not any(kw.arg == arg for kw in call.keywords): + call.keywords.append(ast.keyword(arg=arg, value=value)) ast.fix_missing_locations(tree) - return ast.unparse(tree) + return header + ast.unparse(tree) + + +def add_extra_to_base_model(code: str) -> str: + "Adds `extra='forbid'` to the ConfigDict in BaseModel if not already present." + return _set_base_model_config_kwarg(code, "extra", ast.Constant("forbid")) + + +def add_defer_build_to_base_model(code: str) -> str: + "Adds `defer_build=True` to the ConfigDict in BaseModel if not already present." + return _set_base_model_config_kwarg(code, "defer_build", ast.Constant(True)) diff --git a/docs/03-reference/01-configuration.md b/docs/03-reference/01-configuration.md index 69666c8b..ca6f46d5 100644 --- a/docs/03-reference/01-configuration.md +++ b/docs/03-reference/01-configuration.md @@ -50,6 +50,7 @@ Exactly one of the following parameters is required - they are mutually exclusiv - `files_to_include` (defaults to `[]`) - list of files which will be copied into generated package - `plugins` (defaults to `[]`) - list of plugins to use during generation - `enable_custom_operations` (defaults to `false`) - enables building custom operations. Generates additional files that contains all the classes and methods for generation. +- `defer_model_build` (defaults to `false`) - defers building of generated Pydantic models until they are first used. Sets `defer_build=True` on the generated `BaseModel` and skips the eager `model_rebuild()` calls, so importing the generated package is much faster for large schemas. - `include_typename` (defaults to `true`) - a flag that specifies whether to include the `__typename` field in generated models - `ignore_extra_fields` (defaults to `true`) - when `true`, generated models ignore extra fields returned by the server; set to `false` to add `extra="forbid"` to the base model so unexpected fields raise a validation error - `default_optional_fields_to_none` (defaults to `false`) - when `true`, optional fields in generated models default to `None` instead of being required keyword arguments diff --git a/pyproject.toml b/pyproject.toml index 8824691a..e805c0b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,9 +77,10 @@ features = ["dev", "types", "test", "subscriptions", "opentelemetry"] path = ".venv" [tool.hatch.envs.default.scripts] -lint = ["hatch fmt --check", "hatch run types:check"] +lint = ["hatch check code", "hatch check fmt", "hatch run types:check"] check = [ - "hatch fmt", + "hatch check code --fix", + "hatch check fmt --fix", "hatch test -a -p", "hatch test --cover", "hatch run types:check", @@ -93,10 +94,44 @@ release-notes = "git cliff --latest --strip all" [tool.hatch.envs.types.scripts] check = ["ty check"] +## Static analysis environments +# +# `config-path` stops hatch from injecting its own bundled ruff defaults, so the +# `[tool.ruff]` config below is used verbatim, at the ruff version pinned here. +# `hatch-static-analysis` backs the deprecated `hatch fmt`; it is pinned as well +# so that running it does not reformat the tree with a different ruff version. + +[tool.hatch.envs.hatch-check-code] +config-path = "none" +dependencies = ["ruff>=0.15.0,<0.16.0"] + +[tool.hatch.envs.hatch-check-fmt] +config-path = "none" +dependencies = ["ruff>=0.15.0,<0.16.0"] + [tool.hatch.envs.hatch-static-analysis] config-path = "none" dependencies = ["ruff>=0.15.0,<0.16.0"] +## Type checking environment +# +# `hatch check types` runs pyrefly by default. Point it at ty instead, so that it +# and `hatch run types:check` agree and `# ty: ignore[...]` stays the only +# suppression syntax. `config-path` stops hatch writing a pyrefly.toml we ignore. + +[tool.hatch.envs.hatch-check-types] +python = "3.10" +config-path = "none" +features = ["types", "test", "subscriptions", "opentelemetry"] +dependencies = [] + +[tool.hatch.envs.hatch-check-types.scripts] +check = "ty check {args}" +# ty has no summary-statistics or type-coverage report, so `--summarize` adds +# nothing and `--cover` fails loudly rather than as `coverage: command not found`. +check-summarize = "ty check {args}" +coverage = "python -c \"raise SystemExit('ty has no type-coverage report, so --cover is unsupported')\"" + ## Test environments [tool.hatch.envs.hatch-test] @@ -201,3 +236,9 @@ mark-parentheses = false [tool.pytest.ini_options] testpaths = ["tests"] +# Timing guards compare wall-clock ratios, so they are too noisy for a shared CI +# runner. A `-m` on the command line replaces this one: `hatch test -- -m performance`. +addopts = "-m 'not performance'" +markers = [ + "performance: timing guards for generated code (may be slow; run with -m performance)", +] diff --git a/tests/client_generators/test_client_generator.py b/tests/client_generators/test_client_generator.py index 25e80a80..b5010467 100644 --- a/tests/client_generators/test_client_generator.py +++ b/tests/client_generators/test_client_generator.py @@ -1090,3 +1090,28 @@ def test_add_method_generates_correct_method_body_for_shadowed_data_variable( method_def = class_def.body[0] assert isinstance(method_def, ast.FunctionDef) assert compare_ast(method_def.body, expected_method_body) + + +def test_generate_does_not_put_shared_import_constants_into_generated_module( + async_base_client_import, +): + """Generated modules must not alias the module-level import constants. + + Plugins rewrite imports in place (`ClientForwardRefsPlugin` strips names off + them). When the constant itself was inserted into the module, that mutation + leaked into every later generation in the same process, eventually emitting a + client whose `UnsetType` annotation had no matching import. + """ + generator = ClientGenerator( + base_client_import=async_base_client_import, + arguments_generator=ArgumentsGenerator(schema=GraphQLSchema()), + ) + + module = generator.generate() + + for node in module.body: + if isinstance(node, ast.ImportFrom): + node.names = [] + + assert [alias.name for alias in UNSET_IMPORT.names] == ["UNSET", "UnsetType"] + assert [alias.name for alias in UPLOAD_IMPORT.names] == ["Upload"] diff --git a/tests/main/clients/base_model_options_no_upload/pyproject.toml b/tests/main/clients/base_model_options_no_upload/pyproject.toml new file mode 100644 index 00000000..2db7a3d5 --- /dev/null +++ b/tests/main/clients/base_model_options_no_upload/pyproject.toml @@ -0,0 +1,8 @@ +[tool.ariadne-codegen] +schema_path = "schema.graphql" +queries_path = "queries.graphql" +target_package_name = "base_model_options_no_upload_client" +include_comments = "none" +multipart_uploads = false +defer_model_build = true +ignore_extra_fields = false diff --git a/tests/main/clients/base_model_options_no_upload/queries.graphql b/tests/main/clients/base_model_options_no_upload/queries.graphql new file mode 100644 index 00000000..509715cd --- /dev/null +++ b/tests/main/clients/base_model_options_no_upload/queries.graphql @@ -0,0 +1,13 @@ +query GetUser($id: ID!) { + user(id: $id) { + id + name + } +} + +mutation CreateUser($input: CreateUserInput!) { + createUser(input: $input) { + id + name + } +} diff --git a/tests/main/clients/base_model_options_no_upload/schema.graphql b/tests/main/clients/base_model_options_no_upload/schema.graphql new file mode 100644 index 00000000..9cef190f --- /dev/null +++ b/tests/main/clients/base_model_options_no_upload/schema.graphql @@ -0,0 +1,16 @@ +type Query { + user(id: ID!): User +} + +type Mutation { + createUser(input: CreateUserInput!): User! +} + +type User { + id: ID! + name: String! +} + +input CreateUserInput { + name: String! +} diff --git a/tests/main/clients/defer_model_build/expected_client/__init__.py b/tests/main/clients/defer_model_build/expected_client/__init__.py new file mode 100644 index 00000000..2bfe0202 --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/__init__.py @@ -0,0 +1,34 @@ +from .async_base_client import AsyncBaseClient +from .base_model import BaseModel, Upload +from .client import Client +from .exceptions import ( + GraphQLClientError, + GraphQLClientGraphQLError, + GraphQLClientGraphQLMultiError, + GraphQLClientHttpError, + GraphQLClientInvalidResponseError, +) +from .fragments import UserFields, UserFieldsFriends +from .get_user import GetUser, GetUserUser, GetUserUserFriends +from .input_types import UserFilterInput +from .list_users import ListUsers, ListUsersUsers + +__all__ = [ + "AsyncBaseClient", + "BaseModel", + "Client", + "GetUser", + "GetUserUser", + "GetUserUserFriends", + "GraphQLClientError", + "GraphQLClientGraphQLError", + "GraphQLClientGraphQLMultiError", + "GraphQLClientHttpError", + "GraphQLClientInvalidResponseError", + "ListUsers", + "ListUsersUsers", + "Upload", + "UserFields", + "UserFieldsFriends", + "UserFilterInput", +] diff --git a/tests/main/clients/defer_model_build/expected_client/async_base_client.py b/tests/main/clients/defer_model_build/expected_client/async_base_client.py new file mode 100644 index 00000000..28d78802 --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/async_base_client.py @@ -0,0 +1,406 @@ +import asyncio +import enum +import json +from collections.abc import AsyncIterator +from typing import IO, Any, Optional, Protocol, TypeVar, cast +from uuid import uuid4 + +import httpx +from pydantic import BaseModel +from pydantic_core import to_jsonable_python + +from .base_model import UNSET, Upload +from .exceptions import ( + GraphQLClientError, + GraphQLClientGraphQLMultiError, + GraphQLClientHttpError, + GraphQLClientInvalidMessageFormat, + GraphQLClientInvalidResponseError, +) + +try: + from websockets import ( # type: ignore[import-not-found,unused-ignore] + ClientConnection, + ) + from websockets import ( # type: ignore[import-not-found,unused-ignore] + connect as ws_connect, + ) + from websockets.typing import ( # type: ignore[import-not-found,unused-ignore] + Data, + Origin, + Subprotocol, + ) +except ImportError: + from contextlib import asynccontextmanager + + @asynccontextmanager # type: ignore + async def ws_connect(*args, **kwargs): + raise NotImplementedError("Subscriptions require 'websockets' package.") + yield + + ClientConnection = Any # ty: ignore[invalid-assignment] + Data = Any # ty: ignore[invalid-assignment] + Origin = Any # ty: ignore[invalid-assignment] + + def Subprotocol(*args, **kwargs): # type: ignore # noqa: N802, N803 + raise NotImplementedError("Subscriptions require 'websockets' package.") + + +class Response(Protocol): + status_code: int + + def json(self, **kwargs: Any) -> Any: ... + + +class HttpClient(Protocol): + async def post( + self, + url: Any | str, + json: Any | None = None, + data: Any | None = None, + files: Any | None = None, + headers: Any | None = None, + **kwargs: Any, + ) -> Response: ... + + async def aclose(self) -> None: ... + + +Self = TypeVar("Self", bound="AsyncBaseClient") + +GRAPHQL_TRANSPORT_WS = "graphql-transport-ws" + + +class GraphQLTransportWSMessageType(str, enum.Enum): + CONNECTION_INIT = "connection_init" + CONNECTION_ACK = "connection_ack" + PING = "ping" + PONG = "pong" + SUBSCRIBE = "subscribe" + NEXT = "next" + ERROR = "error" + COMPLETE = "complete" + + +class AsyncBaseClient: + def __init__( + self, + url: str = "", + headers: Optional[dict[str, str]] = None, + http_client: Optional[HttpClient] = None, + ws_url: str = "", + ws_headers: Optional[dict[str, Any]] = None, + ws_origin: Optional[str] = None, + ws_connection_init_payload: Optional[dict[str, Any]] = None, + ) -> None: + self.url = url + self.headers = headers + self.http_client = ( + http_client if http_client else httpx.AsyncClient(headers=headers) + ) + + self.ws_url = ws_url + self.ws_headers = ws_headers or {} + self.ws_origin = Origin(ws_origin) if ws_origin else None + self.ws_connection_init_payload = ws_connection_init_payload + + async def __aenter__(self: Self) -> Self: + return self + + async def __aexit__( + self, + exc_type: object, + exc_val: object, + exc_tb: object, + ) -> None: + await self.http_client.aclose() + + async def execute( + self, + query: str, + operation_name: Optional[str] = None, + variables: Optional[dict[str, Any]] = None, + **kwargs: Any, + ) -> Response: + processed_variables, files, files_map = self._process_variables(variables) + + if files and files_map: + return await self._execute_multipart( + query=query, + operation_name=operation_name, + variables=processed_variables, + files=files, + files_map=files_map, + **kwargs, + ) + + return await self._execute_json( + query=query, + operation_name=operation_name, + variables=processed_variables, + **kwargs, + ) + + def get_data(self, response: Response) -> dict[str, Any]: + if not (200 <= response.status_code <= 299): + raise GraphQLClientHttpError( + status_code=response.status_code, response=response + ) + + try: + response_json = response.json() + except ValueError as exc: + raise GraphQLClientInvalidResponseError(response=response) from exc + + if (not isinstance(response_json, dict)) or ( + "data" not in response_json and "errors" not in response_json + ): + raise GraphQLClientInvalidResponseError(response=response) + + data = response_json.get("data") + errors = response_json.get("errors") + + if errors: + raise GraphQLClientGraphQLMultiError.from_errors_dicts( + errors_dicts=errors, data=data + ) + + return cast(dict[str, Any], data) + + async def execute_ws( + self, + query: str, + operation_name: Optional[str] = None, + variables: Optional[dict[str, Any]] = None, + **kwargs: Any, + ) -> AsyncIterator[dict[str, Any]]: + headers = self.ws_headers.copy() + headers.update(kwargs.pop("additional_headers", {})) + + merged_kwargs: dict[str, Any] = {"origin": self.ws_origin} + merged_kwargs.update(kwargs) + merged_kwargs["additional_headers"] = headers + + operation_id = str(uuid4()) + async with ws_connect( + self.ws_url, + subprotocols=[Subprotocol(GRAPHQL_TRANSPORT_WS)], + **merged_kwargs, + ) as websocket: + await self._send_connection_init(websocket) + # Wait for connection_ack; some servers (e.g. Hasura) send ping before + # connection_ack, so we loop and handle pings until we get ack. + try: + await asyncio.wait_for( + self._wait_for_connection_ack(websocket), + timeout=5.0, + ) + except asyncio.TimeoutError as exc: + raise GraphQLClientError( + "Connection ack not received within 5 seconds" + ) from exc + await self._send_subscribe( + websocket, + operation_id=operation_id, + query=query, + operation_name=operation_name, + variables=variables, + ) + + async for message in websocket: + data = await self._handle_ws_message(message, websocket) + if data and "connection_ack" not in data: + yield data + + def _process_variables( + self, variables: Optional[dict[str, Any]] + ) -> tuple[ + dict[str, Any], dict[str, tuple[str, IO[bytes], str]], dict[str, list[str]] + ]: + if not variables: + return {}, {}, {} + + serializable_variables = self._convert_dict_to_json_serializable(variables) + return self._get_files_from_variables(serializable_variables) + + def _convert_dict_to_json_serializable( + self, dict_: dict[str, Any] + ) -> dict[str, Any]: + return { + key: self._convert_value(value) + for key, value in dict_.items() + if value is not UNSET + } + + def _convert_value(self, value: Any) -> Any: + if isinstance(value, BaseModel): + return value.model_dump(by_alias=True, exclude_unset=True) + if isinstance(value, list): + return [self._convert_value(item) for item in value] + return value + + def _get_files_from_variables( + self, variables: dict[str, Any] + ) -> tuple[ + dict[str, Any], dict[str, tuple[str, IO[bytes], str]], dict[str, list[str]] + ]: + files_map: dict[str, list[str]] = {} + files_list: list[Upload] = [] + + def separate_files(path: str, obj: Any) -> Any: + if isinstance(obj, list): + nulled_list = [] + for index, value in enumerate(obj): + value = separate_files(f"{path}.{index}", value) + nulled_list.append(value) + return nulled_list + + if isinstance(obj, dict): + nulled_dict = {} + for key, value in obj.items(): + value = separate_files(f"{path}.{key}", value) + nulled_dict[key] = value + return nulled_dict + + if isinstance(obj, Upload): + if obj in files_list: + file_index = files_list.index(obj) + files_map[str(file_index)].append(path) + else: + file_index = len(files_list) + files_list.append(obj) + files_map[str(file_index)] = [path] + return None + + return obj + + nulled_variables = separate_files("variables", variables) + files: dict[str, tuple[str, IO[bytes], str]] = { + str(i): (file_.filename, cast(IO[bytes], file_.content), file_.content_type) + for i, file_ in enumerate(files_list) + } + return nulled_variables, files, files_map + + async def _execute_multipart( + self, + query: str, + operation_name: Optional[str], + variables: dict[str, Any], + files: dict[str, tuple[str, IO[bytes], str]], + files_map: dict[str, list[str]], + **kwargs: Any, + ) -> Response: + data = { + "operations": json.dumps( + { + "query": query, + "operationName": operation_name, + "variables": variables, + }, + default=to_jsonable_python, + ), + "map": json.dumps(files_map, default=to_jsonable_python), + } + + return await self.http_client.post( + url=self.url, data=data, files=files, **kwargs + ) + + async def _execute_json( + self, + query: str, + operation_name: Optional[str], + variables: dict[str, Any], + **kwargs: Any, + ) -> Response: + return await self.http_client.post( + url=self.url, + json=to_jsonable_python( + { + "query": query, + "operationName": operation_name, + "variables": variables, + } + ), + **kwargs, + ) + + async def _send_connection_init(self, websocket: ClientConnection) -> None: + payload: dict[str, Any] = { + "type": GraphQLTransportWSMessageType.CONNECTION_INIT.value + } + if self.ws_connection_init_payload: + payload["payload"] = self.ws_connection_init_payload + await websocket.send(json.dumps(payload)) + + async def _wait_for_connection_ack(self, websocket: ClientConnection) -> None: + """Read messages until connection_ack; handle ping/pong in between.""" + async for message in websocket: + data = await self._handle_ws_message(message, websocket) + if data is not None and "connection_ack" in data: + return + + async def _send_subscribe( + self, + websocket: ClientConnection, + operation_id: str, + query: str, + operation_name: Optional[str] = None, + variables: Optional[dict[str, Any]] = None, + ) -> None: + payload_inner: dict[str, Any] = { + "query": query, + "operationName": operation_name, + } + if variables: + payload_inner["variables"] = self._convert_dict_to_json_serializable( + variables + ) + payload: dict[str, Any] = { + "id": operation_id, + "type": GraphQLTransportWSMessageType.SUBSCRIBE.value, + "payload": payload_inner, + } + await websocket.send(json.dumps(payload)) + + async def _handle_ws_message( + self, + message: Data, + websocket: ClientConnection, + expected_type: Optional[GraphQLTransportWSMessageType] = None, + ) -> Optional[dict[str, Any]]: + try: + message_dict = json.loads(message) + except json.JSONDecodeError as exc: + raise GraphQLClientInvalidMessageFormat(message=message) from exc + + type_ = message_dict.get("type") + payload = message_dict.get("payload", {}) + + if not type_ or type_ not in {t.value for t in GraphQLTransportWSMessageType}: + raise GraphQLClientInvalidMessageFormat(message=message) + + if expected_type and expected_type != type_: + raise GraphQLClientInvalidMessageFormat( + f"Invalid message received. Expected: {expected_type.value}" + ) + + if type_ == GraphQLTransportWSMessageType.NEXT: + if "data" not in payload: + raise GraphQLClientInvalidMessageFormat(message=message) + return cast(dict[str, Any], payload["data"]) + + if type_ == GraphQLTransportWSMessageType.COMPLETE: + await websocket.close() + elif type_ == GraphQLTransportWSMessageType.PING: + await websocket.send( + json.dumps({"type": GraphQLTransportWSMessageType.PONG.value}) + ) + elif type_ == GraphQLTransportWSMessageType.ERROR: + raise GraphQLClientGraphQLMultiError.from_errors_dicts( + errors_dicts=payload, data=message_dict + ) + elif type_ == GraphQLTransportWSMessageType.CONNECTION_ACK: + return {"connection_ack": True} + + return None diff --git a/tests/main/clients/defer_model_build/expected_client/base_model.py b/tests/main/clients/defer_model_build/expected_client/base_model.py new file mode 100644 index 00000000..36fee87a --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/base_model.py @@ -0,0 +1,29 @@ +from io import IOBase + +from pydantic import BaseModel as PydanticBaseModel +from pydantic import ConfigDict + + +class UnsetType: + def __bool__(self) -> bool: + return False + + +UNSET = UnsetType() + + +class BaseModel(PydanticBaseModel): + model_config = ConfigDict( + populate_by_name=True, + validate_assignment=True, + arbitrary_types_allowed=True, + protected_namespaces=(), + defer_build=True, + ) + + +class Upload: + def __init__(self, filename: str, content: IOBase, content_type: str): + self.filename = filename + self.content = content + self.content_type = content_type diff --git a/tests/main/clients/defer_model_build/expected_client/client.py b/tests/main/clients/defer_model_build/expected_client/client.py new file mode 100644 index 00000000..6c8c87a6 --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/client.py @@ -0,0 +1,61 @@ +from typing import Any, Optional, Union + +from .async_base_client import AsyncBaseClient +from .base_model import UNSET, UnsetType +from .get_user import GetUser +from .input_types import UserFilterInput +from .list_users import ListUsers + + +def gql(q: str) -> str: + return q + + +class Client(AsyncBaseClient): + async def get_user(self, id: str, **kwargs: Any) -> GetUser: + query = gql(""" + query GetUser($id: ID!) { + user(id: $id) { + id + name + friends { + id + name + } + } + } + """) + variables: dict[str, object] = {"id": id} + response = await self.execute( + query=query, operation_name="GetUser", variables=variables, **kwargs + ) + data = self.get_data(response) + return GetUser.model_validate(data) + + async def list_users( + self, + filter_: Union[Optional[UserFilterInput], UnsetType] = UNSET, + **kwargs: Any, + ) -> ListUsers: + query = gql(""" + query ListUsers($filter: UserFilterInput) { + users(filter: $filter) { + ...UserFields + } + } + + fragment UserFields on User { + id + name + friends { + id + name + } + } + """) + variables: dict[str, object] = {"filter": filter_} + response = await self.execute( + query=query, operation_name="ListUsers", variables=variables, **kwargs + ) + data = self.get_data(response) + return ListUsers.model_validate(data) diff --git a/tests/main/clients/defer_model_build/expected_client/enums.py b/tests/main/clients/defer_model_build/expected_client/enums.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/main/clients/defer_model_build/expected_client/exceptions.py b/tests/main/clients/defer_model_build/expected_client/exceptions.py new file mode 100644 index 00000000..885d6663 --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/exceptions.py @@ -0,0 +1,85 @@ +from typing import Any, Optional, Protocol, Union + + +class Response(Protocol): + status_code: int + + +class GraphQLClientError(Exception): + """Base exception.""" + + +class GraphQLClientHttpError(GraphQLClientError): + def __init__(self, status_code: int, response: Response) -> None: + self.status_code = status_code + self.response = response + + def __str__(self) -> str: + return f"HTTP status code: {self.status_code}" + + +class GraphQLClientInvalidResponseError(GraphQLClientError): + def __init__(self, response: Response) -> None: + self.response = response + + def __str__(self) -> str: + return "Invalid response format." + + +class GraphQLClientGraphQLError(GraphQLClientError): + def __init__( + self, + message: str, + locations: Optional[list[dict[str, int]]] = None, + path: Optional[list[str]] = None, + extensions: Optional[dict[str, object]] = None, + original: Optional[dict[str, object]] = None, + ): + self.message = message + self.locations = locations + self.path = path + self.extensions = extensions + self.original = original + + def __str__(self) -> str: + return self.message + + @classmethod + def from_dict(cls, error: dict[str, Any]) -> "GraphQLClientGraphQLError": + return cls( + message=error["message"], + locations=error.get("locations"), + path=error.get("path"), + extensions=error.get("extensions"), + original=error, + ) + + +class GraphQLClientGraphQLMultiError(GraphQLClientError): + def __init__( + self, + errors: list[GraphQLClientGraphQLError], + data: Optional[dict[str, Any]] = None, + ): + self.errors = errors + self.data = data + + def __str__(self) -> str: + return "; ".join(str(e) for e in self.errors) + + @classmethod + def from_errors_dicts( + cls, errors_dicts: list[dict[str, Any]], data: Optional[dict[str, Any]] = None + ) -> "GraphQLClientGraphQLMultiError": + return cls( + errors=[GraphQLClientGraphQLError.from_dict(e) for e in errors_dicts], + data=data, + ) + + +class GraphQLClientInvalidMessageFormat(GraphQLClientError): # noqa: N818 + def __init__(self, message: Union[str, bytes]) -> None: + self.message = message + + def __str__(self) -> str: + return "Invalid message format." diff --git a/tests/main/clients/defer_model_build/expected_client/fragments.py b/tests/main/clients/defer_model_build/expected_client/fragments.py new file mode 100644 index 00000000..9017b28d --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/fragments.py @@ -0,0 +1,12 @@ +from .base_model import BaseModel + + +class UserFields(BaseModel): + id: str + name: str + friends: list["UserFieldsFriends"] + + +class UserFieldsFriends(BaseModel): + id: str + name: str diff --git a/tests/main/clients/defer_model_build/expected_client/get_user.py b/tests/main/clients/defer_model_build/expected_client/get_user.py new file mode 100644 index 00000000..0d9f8533 --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/get_user.py @@ -0,0 +1,18 @@ +from typing import Optional + +from .base_model import BaseModel + + +class GetUser(BaseModel): + user: Optional["GetUserUser"] + + +class GetUserUser(BaseModel): + id: str + name: str + friends: list["GetUserUserFriends"] + + +class GetUserUserFriends(BaseModel): + id: str + name: str diff --git a/tests/main/clients/defer_model_build/expected_client/input_types.py b/tests/main/clients/defer_model_build/expected_client/input_types.py new file mode 100644 index 00000000..b67e99b7 --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/input_types.py @@ -0,0 +1,9 @@ +from typing import Optional + +from .base_model import BaseModel + + +class UserFilterInput(BaseModel): + name: Optional[str] = None + manager: Optional["UserFilterInput"] = None + tags: Optional[list[str]] = None diff --git a/tests/main/clients/defer_model_build/expected_client/list_users.py b/tests/main/clients/defer_model_build/expected_client/list_users.py new file mode 100644 index 00000000..25e4e816 --- /dev/null +++ b/tests/main/clients/defer_model_build/expected_client/list_users.py @@ -0,0 +1,10 @@ +from .base_model import BaseModel +from .fragments import UserFields + + +class ListUsers(BaseModel): + users: list["ListUsersUsers"] + + +class ListUsersUsers(UserFields): + pass diff --git a/tests/main/clients/defer_model_build/pyproject.toml b/tests/main/clients/defer_model_build/pyproject.toml new file mode 100644 index 00000000..aa6a1f9d --- /dev/null +++ b/tests/main/clients/defer_model_build/pyproject.toml @@ -0,0 +1,6 @@ +[tool.ariadne-codegen] +schema_path = "schema.graphql" +queries_path = "queries.graphql" +include_comments = "none" +target_package_name = "defer_model_build_client" +defer_model_build = true diff --git a/tests/main/clients/defer_model_build/queries.graphql b/tests/main/clients/defer_model_build/queries.graphql new file mode 100644 index 00000000..bf6c018d --- /dev/null +++ b/tests/main/clients/defer_model_build/queries.graphql @@ -0,0 +1,25 @@ +fragment UserFields on User { + id + name + friends { + id + name + } +} + +query GetUser($id: ID!) { + user(id: $id) { + id + name + friends { + id + name + } + } +} + +query ListUsers($filter: UserFilterInput) { + users(filter: $filter) { + ...UserFields + } +} diff --git a/tests/main/clients/defer_model_build/schema.graphql b/tests/main/clients/defer_model_build/schema.graphql new file mode 100644 index 00000000..23f971b0 --- /dev/null +++ b/tests/main/clients/defer_model_build/schema.graphql @@ -0,0 +1,21 @@ +schema { + query: Query +} + +type Query { + user(id: ID!): User + users(filter: UserFilterInput): [User!]! +} + +type User { + id: ID! + name: String! + friends: [User!]! + manager: User +} + +input UserFilterInput { + name: String + manager: UserFilterInput + tags: [String!] +} diff --git a/tests/main/test_main.py b/tests/main/test_main.py index ca48044d..828c4130 100644 --- a/tests/main/test_main.py +++ b/tests/main/test_main.py @@ -254,6 +254,17 @@ def test_main_shows_version(): "expected_client", CLIENTS_PATH / "no_multipart_upload" / "expected_client", ), + ( + ( + CLIENTS_PATH / "defer_model_build" / "pyproject.toml", + ( + CLIENTS_PATH / "defer_model_build" / "queries.graphql", + CLIENTS_PATH / "defer_model_build" / "schema.graphql", + ), + ), + "defer_model_build_client", + CLIENTS_PATH / "defer_model_build" / "expected_client", + ), ], indirect=["project_dir"], ) @@ -268,6 +279,86 @@ def test_main_generates_correct_package( assert_the_same_files_in_directories(package_path, expected_package_path) +@pytest.mark.parametrize( + "project_dir, package_name", + [ + ( + ( + CLIENTS_PATH / "defer_model_build" / "pyproject.toml", + ( + CLIENTS_PATH / "defer_model_build" / "queries.graphql", + CLIENTS_PATH / "defer_model_build" / "schema.graphql", + ), + ), + "defer_model_build_client", + ), + ], + indirect=["project_dir"], +) +def test_main_with_defer_model_build_omits_model_rebuild_calls( + project_dir, package_name +): + """With ``defer_model_build = true`` the generated package must not contain + any eager ``model_rebuild()`` calls, and ``BaseModel`` must enable + ``defer_build`` so Pydantic builds each model schema lazily on first use. + """ + result = CliRunner().invoke(main) + + assert result.exit_code == 0 + package_path = project_dir / package_name + assert package_path.is_dir() + + for module_path in package_path.glob("*.py"): + assert "model_rebuild()" not in module_path.read_text(), ( + f"unexpected model_rebuild() in {module_path.name}" + ) + + base_model_code = (package_path / "base_model.py").read_text() + assert "defer_build=True" in base_model_code + + # Forward references must still be generated (otherwise there would be + # nothing to defer). This guards against accidentally dropping them. + assert ( + 'Optional["UserFilterInput"]' in (package_path / "input_types.py").read_text() + ) + + +@pytest.mark.parametrize( + "project_dir, package_name", + [ + ( + ( + CLIENTS_PATH / "base_model_options_no_upload" / "pyproject.toml", + ( + CLIENTS_PATH / "base_model_options_no_upload" / "queries.graphql", + CLIENTS_PATH / "base_model_options_no_upload" / "schema.graphql", + ), + ), + "base_model_options_no_upload_client", + ), + ], + indirect=["project_dir"], +) +def test_main_applies_base_model_config_with_multipart_uploads_disabled( + project_dir, package_name +): + """The base-model rewrites must apply even when ``multipart_uploads = false``. + + With uploads disabled the base model is copied from ``base_model_no_upload.py`` + rather than ``base_model.py``, so keying the rewrite off the source file name + silently skipped it, shipping a ``BaseModel`` without the requested + ``defer_build``/``extra`` config while still reporting the options as enabled. + The rewrite is keyed off the base-model source path, so it must fire in both + upload modes. + """ + result = CliRunner().invoke(main) + + assert result.exit_code == 0 + base_model_code = (project_dir / package_name / "base_model.py").read_text() + assert "defer_build=True" in base_model_code + assert 'extra="forbid"' in base_model_code + + @pytest.mark.parametrize( "project_dir, package_name", [ diff --git a/tests/performance/__init__.py b/tests/performance/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/performance/helpers.py b/tests/performance/helpers.py new file mode 100644 index 00000000..b5f9ad09 --- /dev/null +++ b/tests/performance/helpers.py @@ -0,0 +1,90 @@ +"""Shared helpers for the performance guards. + +These tests compare two ways of doing the same thing and assert a ratio, rather +than asserting an absolute duration, so that they stay meaningful on machines of +wildly different speed. Every threshold is set well below the ratio the change +actually delivers, so what they catch is a change of *mechanism* (somebody +re-introducing eager `model_rebuild()` calls, or per-file ruff spawns), not a +slow CI runner. +""" + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +from click.testing import CliRunner + +from ariadne_codegen.main import main + +# Number of fresh-subprocess timings to take per measurement; the smallest is +# used, which is far more stable than a mean when the noise is all one-sided. +TIMING_REPEATS = 3 + + +def report(request: pytest.FixtureRequest, title: str, rows: list[str]) -> None: + """Print a comparison table through the terminal reporter. + + Written straight to the reporter rather than with `print`, so the numbers + show up on a normal run instead of only under `-s`. + """ + reporter = request.config.pluginmanager.get_plugin("terminalreporter") + if reporter is None: + return + reporter.ensure_newline() + reporter.write_line(f" {title}") + for row in rows: + reporter.write_line(f" {row}") + + +def generate_package(project_dir: Path, package_name: str, settings: str = "") -> Path: + """Run codegen in `project_dir` and return the generated package path.""" + (project_dir / "pyproject.toml").write_text( + textwrap.dedent( + f"""\ + [tool.ariadne-codegen] + schema_path = "schema.graphql" + queries_path = "queries.graphql" + include_comments = "none" + target_package_name = "{package_name}" + """ + ) + + settings + ) + + old_cwd = Path.cwd() + os.chdir(project_dir) + try: + result = CliRunner().invoke(main) + assert result.exit_code == 0, result.output + finally: + os.chdir(old_cwd) + + package_path = project_dir / package_name + assert package_path.is_dir() + return package_path + + +def measure_import_seconds(project_dir: Path, package_name: str) -> float: + """Import the package in a fresh interpreter; return the best wall-clock time.""" + script = textwrap.dedent( + f""" + import time + start = time.perf_counter() + import {package_name} # noqa: F401 + print(time.perf_counter() - start) + """ + ) + best = float("inf") + for _ in range(TIMING_REPEATS): + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=project_dir, + capture_output=True, + text=True, + check=True, + ) + best = min(best, float(completed.stdout.strip().splitlines()[-1])) + return best diff --git a/tests/performance/test_codegen_performance.py b/tests/performance/test_codegen_performance.py new file mode 100644 index 00000000..09ada2d6 --- /dev/null +++ b/tests/performance/test_codegen_performance.py @@ -0,0 +1,101 @@ +"""Guard on how many times generation spawns ruff. + +Formatting used to spawn ruff twice per generated file, a spawn count linear in +the size of the package. Ruff is now invoked once per rule-selection group for +the whole package. The test asserts that mechanism, that the spawn count stays +constant as the package grows, rather than wall-clock time, so it does not flake +on a loaded runner. + +Marked `performance`, so it can be run on its own: + + hatch test -- -m performance +""" + +import subprocess +from pathlib import Path +from unittest import mock + +import pytest + +import ariadne_codegen.utils as utils + +from .helpers import generate_package, report + +pytestmark = pytest.mark.performance + +# `ariadne_codegen.utils.subprocess` is the stdlib module, so patching its `run` +# patches it everywhere. Hold on to the real one to wrap with. +_REAL_SUBPROCESS_RUN = subprocess.run + +# Two package sizes far enough apart that a per-file spawn count could not +# coincide between them. +SMALL_OPERATION_COUNT = 5 +LARGE_OPERATION_COUNT = 60 + +# Ruff runs once per rule-selection group (`I` and `I,F401`), for `check --fix` +# and again for `format`. Nothing here rewrites base_model.py, which would add a +# further pair. The bound matters less than the fact that it does not scale. +MAX_RUFF_SPAWNS = 6 + + +def _build_project(project_dir: Path, operation_count: int) -> None: + project_dir.mkdir(parents=True, exist_ok=True) + types = "\n".join(f"type T{i} {{ a: String b: Int }}" for i in range(20)) + query_fields = " ".join( + f"get{i}(id: ID!): T{i % 20}" for i in range(operation_count) + ) + (project_dir / "schema.graphql").write_text( + f"schema {{ query: Query }}\n{types}\ntype Query {{ {query_fields} }}\n" + ) + (project_dir / "queries.graphql").write_text( + "\n".join( + f"query Op{i}($id: ID!) {{ get{i}(id: $id) {{ a b }} }}" + for i in range(operation_count) + ) + ) + + +def _generate_recording_spawns(project_dir: Path) -> tuple[list, Path]: + with mock.patch.object(utils.subprocess, "run", wraps=_REAL_SUBPROCESS_RUN) as spy: + package_path = generate_package(project_dir, "codegen_client") + return spy.call_args_list, package_path + + +def _ruff_spawns(calls: list) -> list: + return [ + call + for call in calls + if len(call[0][0]) > 1 and call[0][0][1] in ("check", "format") + ] + + +def test_ruff_spawn_count_does_not_grow_with_the_package(request, tmp_path): + """The spawn count is what makes generation scale; wall-clock only reflects it.""" + _build_project(tmp_path / "small", SMALL_OPERATION_COUNT) + small_calls, small_package = _generate_recording_spawns(tmp_path / "small") + + _build_project(tmp_path / "large", LARGE_OPERATION_COUNT) + large_calls, large_package = _generate_recording_spawns(tmp_path / "large") + + small_spawns = len(_ruff_spawns(small_calls)) + large_spawns = len(_ruff_spawns(large_calls)) + small_files = len(list(small_package.glob("*.py"))) + large_files = len(list(large_package.glob("*.py"))) + + report( + request, + "ruff subprocess spawns per generated package", + [ + f"{small_files:3d} files -> {small_spawns} spawns", + f"{large_files:3d} files -> {large_spawns} spawns" + f" (per-file would be {2 * large_files})", + ], + ) + + assert large_files > small_files * 2, "the two packages must differ in size" + assert large_spawns == small_spawns, ( + f"ruff spawn count scales with the package: " + f"{small_files} files -> {small_spawns} spawns, " + f"{large_files} files -> {large_spawns} spawns" + ) + assert large_spawns <= MAX_RUFF_SPAWNS diff --git a/tests/performance/test_defer_model_build_performance.py b/tests/performance/test_defer_model_build_performance.py new file mode 100644 index 00000000..3f01ec3c --- /dev/null +++ b/tests/performance/test_defer_model_build_performance.py @@ -0,0 +1,154 @@ +"""Performance guards for the ``defer_model_build`` setting. + +These tests generate a large, densely interconnected client both with and +without ``defer_model_build`` and compare how long the generated package takes +to *import*. With eager model building Pydantic constructs every model's +core-schema at import time, and because each model inlines the schemas of the +types it references that cost grows super-linearly with a big schema. Deferring +the build moves that work to first use, so importing the package becomes +dramatically cheaper. + +The assertions are deliberately ratio-based (and use a generous margin) so they +guard against regressions - i.e. someone accidentally re-introducing eager +``model_rebuild()`` calls - without being flaky across machines. + +Marked ``performance`` so they can be selected/deselected explicitly: + + hatch test -- -m performance +""" + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +from click.testing import CliRunner + +from ariadne_codegen.main import main + +# Schema size. Large enough that eager building is clearly more expensive than +# deferred, small enough that generating it twice stays fast in CI. The build +# cost is dominated by the number of fields per model, so we use many models +# each with a wide set of fields. Every model keeps a single self forward +# reference so the eager variant still emits ``model_rebuild()`` calls. +NUMBER_OF_MODELS = 300 +FIELDS_PER_MODEL = 20 +NUMBER_OF_HUBS = 6 + +# Deferred import must be at least this many times faster than eager. On this +# schema the real ratio is ~2x; if eager ``model_rebuild()`` calls are ever +# re-introduced for the deferred variant the ratio collapses to ~1x, which this +# bound catches while staying stable on slow/noisy CI runners. +MIN_SPEEDUP = 1.5 + +# Number of fresh-subprocess import timings to take per package; the smallest is +# used to reduce the impact of scheduler noise. +TIMING_REPEATS = 3 + + +def _build_schema() -> str: + lines = [ + "schema { query: Query }", + "type Query { ping(filter: Model0): String }", + ] + for hub in range(NUMBER_OF_HUBS): + lines.append( + f"input Hub{hub} {{ id: Int name: String value: Float flag: Boolean }}" + ) + for i in range(NUMBER_OF_MODELS): + fields = [f"f{k}: String" for k in range(FIELDS_PER_MODEL)] + fields += [f"i{k}: Int" for k in range(FIELDS_PER_MODEL // 2)] + fields += [f"rel{hub}: Hub{hub}" for hub in range(NUMBER_OF_HUBS)] + fields.append(f"selfRef: Model{i}") + fields.append(f"children: [Model{i}!]") + lines.append("input Model{} {{ {} }}".format(i, " ".join(fields))) + return "\n".join(lines) + + +QUERIES = "query Ping {\n ping\n}\n" + + +def _pyproject(package_name: str, defer: bool) -> str: + return textwrap.dedent( + f"""\ + [tool.ariadne-codegen] + schema_path = "schema.graphql" + queries_path = "queries.graphql" + include_comments = "none" + target_package_name = "{package_name}" + defer_model_build = {"true" if defer else "false"} + """ + ) + + +def _generate_package(project_dir: Path, package_name: str, defer: bool) -> Path: + project_dir.mkdir(parents=True, exist_ok=True) + (project_dir / "schema.graphql").write_text(_build_schema()) + (project_dir / "queries.graphql").write_text(QUERIES) + (project_dir / "pyproject.toml").write_text(_pyproject(package_name, defer)) + + old_cwd = Path.cwd() + os.chdir(project_dir) + try: + result = CliRunner().invoke(main) + assert result.exit_code == 0, result.output + finally: + os.chdir(old_cwd) + + package_path = project_dir / package_name + assert package_path.is_dir() + return package_path + + +def _measure_import_seconds(project_dir: Path, package_name: str) -> float: + """Import the generated package in a fresh interpreter and return the best + (smallest) wall-clock time across a few repeats.""" + script = textwrap.dedent( + f""" + import time + start = time.perf_counter() + import {package_name} # noqa: F401 + print(time.perf_counter() - start) + """ + ) + best = float("inf") + for _ in range(TIMING_REPEATS): + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=project_dir, + capture_output=True, + text=True, + check=True, + ) + best = min(best, float(completed.stdout.strip().splitlines()[-1])) + return best + + +@pytest.mark.performance +def test_defer_model_build_speeds_up_import(tmp_path): + eager_dir = tmp_path / "eager" + deferred_dir = tmp_path / "deferred" + + eager_package = _generate_package(eager_dir, "eager_client", defer=False) + deferred_package = _generate_package(deferred_dir, "deferred_client", defer=True) + + # Sanity: the eager package builds models at import time (model_rebuild), + # the deferred one must not. + assert any( + "model_rebuild()" in p.read_text() for p in eager_package.glob("*.py") + ), "expected eager package to contain model_rebuild() calls" + assert all( + "model_rebuild()" not in p.read_text() for p in deferred_package.glob("*.py") + ), "deferred package must not contain model_rebuild() calls" + + eager_seconds = _measure_import_seconds(eager_dir, "eager_client") + deferred_seconds = _measure_import_seconds(deferred_dir, "deferred_client") + + speedup = eager_seconds / deferred_seconds + assert speedup >= MIN_SPEEDUP, ( + f"defer_model_build import not fast enough: " + f"eager={eager_seconds:.3f}s deferred={deferred_seconds:.3f}s " + f"speedup={speedup:.2f}x (expected >= {MIN_SPEEDUP}x)" + ) diff --git a/tests/test_settings.py b/tests/test_settings.py index e46ecdc7..dc6646f0 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -833,3 +833,33 @@ def test_graphql_schema_settings_used_settings_message_includes_schema_paths(): result = settings.used_settings_message assert "pkg.SCHEMA_DIR" in result + + +def test_client_settings_defer_model_build_defaults_to_false(tmp_path): + schema_path = tmp_path / "schema.graphql" + schema_path.touch() + queries_path = tmp_path / "queries.graphql" + queries_path.touch() + + settings = ClientSettings( + schema_path=schema_path.as_posix(), + queries_path=queries_path.as_posix(), + ) + + assert settings.defer_model_build is False + + +def test_client_settings_defer_model_build_can_be_set_to_true(tmp_path): + schema_path = tmp_path / "schema.graphql" + schema_path.touch() + queries_path = tmp_path / "queries.graphql" + queries_path.touch() + + settings = ClientSettings( + schema_path=schema_path.as_posix(), + queries_path=queries_path.as_posix(), + defer_model_build=True, + ) + + assert settings.defer_model_build is True + assert "Deferring Pydantic model builds" in settings.used_settings_message diff --git a/tests/test_utils.py b/tests/test_utils.py index 84e57620..485bbf42 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,14 +1,17 @@ import ast import subprocess +import sys from textwrap import dedent import pytest from ariadne_codegen.utils import ( _format_code, + add_defer_build_to_base_model, add_extra_to_base_model, ast_to_str, convert_to_multiline_string, + format_many, format_multiline_strings, get_variable_indent_size, process_name, @@ -88,21 +91,68 @@ def test_ast_to_str_non_ascii_unicode_round_trip_issue_422(): assert description in generated +def _ruff_calls(spy, subcommand: str) -> list: + """Calls to `ruff `, whatever prefix `_ruff_command` resolved to.""" + return [ + call + for call in spy.call_args_list + if subcommand in call[0][0] and call[0][0].index(subcommand) <= 3 + ] + + def test_format_code_ruff_format_uses_utf8_encoding_issue_422(mocker): """Ensure ruff format stdin/stdout use UTF-8 (mirumee/ariadne-codegen#422).""" spy = mocker.patch("ariadne_codegen.utils.subprocess.run", wraps=subprocess.run) _format_code("x = 1\n") - format_calls = [ - call - for call in spy.call_args_list - if tuple(call[0][0][2:4]) == ("ruff", "format") - ] + format_calls = _ruff_calls(spy, "format") assert format_calls, "expected a ruff format subprocess.run" assert format_calls[-1][1].get("encoding") == "utf-8" +def test_format_code_invokes_ruff_binary_directly(mocker): + """Spawning `python -m ruff` per file dominates generation time.""" + spy = mocker.patch("ariadne_codegen.utils.subprocess.run", wraps=subprocess.run) + + _format_code("x = 1\n") + + for call in spy.call_args_list: + assert call[0][0][:2] != [sys.executable, "-m"], ( + "ruff should be spawned as a binary, not through a Python interpreter" + ) + + +def test_format_code_keeps_empty_result_when_every_import_is_removed(): + """`ruff check --fix` legitimately returns '' here; it must not be discarded.""" + assert _format_code("from enum import Enum\n") == "" + assert _format_code("from enum import Enum\n", remove_unused_imports=False) == ( + "from enum import Enum\n" + ) + + +@pytest.mark.parametrize("remove_unused_imports", [True, False]) +def test_format_many_matches_format_code_per_module(remove_unused_imports): + codes = [ + "import os\nx=1\n", + "from enum import Enum\n", + "from typing import Optional\nclass A:\n\n x: Optional[int]=None\n", + "", + ] + + assert format_many(codes, remove_unused_imports=remove_unused_imports) == [ + _format_code(code, remove_unused_imports=remove_unused_imports) + for code in codes + ] + + +def test_format_many_without_modules_does_not_spawn_ruff(mocker): + spy = mocker.patch("ariadne_codegen.utils.subprocess.run", wraps=subprocess.run) + + assert format_many([]) == [] + assert not spy.call_args_list + + @pytest.mark.parametrize( "name, expected_result", [ @@ -266,3 +316,68 @@ class NotBaseModel: """) result = add_extra_to_base_model(code) assert dedent(result).strip() == expected.strip() + + +def test_adds_defer_build_to_base_model_if_missing(): + code = dedent(""" + class BaseModel: + model_config = ConfigDict(populate_by_name=True) + """) + expected = dedent(""" + class BaseModel: + model_config = ConfigDict(populate_by_name=True, defer_build=True) + """) + result = add_defer_build_to_base_model(code) + assert dedent(result).strip() == expected.strip() + + +def test_adds_defer_build_to_base_model_does_not_overwrite_existing_value(): + code = dedent(""" + class BaseModel: + model_config = ConfigDict(defer_build=False) + """) + expected = dedent(""" + class BaseModel: + model_config = ConfigDict(defer_build=False) + """) + result = add_defer_build_to_base_model(code) + assert dedent(result).strip() == expected.strip() + + +def test_adds_defer_build_to_base_model_leaves_other_classes_untouched(): + code = dedent(""" + class NotBaseModel: + model_config = ConfigDict() + """) + expected = dedent(""" + class NotBaseModel: + model_config = ConfigDict() + """) + result = add_defer_build_to_base_model(code) + assert dedent(result).strip() == expected.strip() + + +@pytest.mark.parametrize( + "rewrite", [add_extra_to_base_model, add_defer_build_to_base_model] +) +def test_rewriting_base_model_keeps_the_generated_file_header(rewrite): + """The rewrites round-trip through `ast.unparse`, which drops comments.""" + code = dedent( + """\ + # Generated by ariadne-codegen + # Source: queries.graphql + + from pydantic import BaseModel as PydanticBaseModel + from pydantic import ConfigDict + + + class BaseModel(PydanticBaseModel): + model_config = ConfigDict(populate_by_name=True) + """ + ) + + result = rewrite(code) + + assert result.startswith( + "# Generated by ariadne-codegen\n# Source: queries.graphql" + ) diff --git a/uv.lock b/uv.lock index b611eef6..1fae9fa1 100644 --- a/uv.lock +++ b/uv.lock @@ -274,7 +274,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -304,11 +304,11 @@ wheels = [ [[package]] name = "graphql-core" -version = "3.2.8" +version = "3.2.11" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/68/c5/36aa96205c3ecbb3d34c7c24189e4553c7ca2ebc7e1dd07432339b980272/graphql_core-3.2.8.tar.gz", hash = "sha256:015457da5d996c924ddf57a43f4e959b0b94fb695b85ed4c29446e508ed65cf3", size = 513181, upload-time = "2026-03-05T19:55:37.332Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4d/90/f2aff026ab4aebd80eb71905106a0885f4cfde85dcf965543f45bed0d9ee/graphql_core-3.2.11.tar.gz", hash = "sha256:e7e156d10beb127cab5c89ff0da71416fc73d27c484a4757d3b2d35633774802", size = 528407, upload-time = "2026-06-05T13:45:22.915Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/86/41/cb887d9afc5dabd78feefe6ccbaf83ff423c206a7a1b7aeeac05120b2125/graphql_core-3.2.8-py3-none-any.whl", hash = "sha256:cbee07bee1b3ed5e531723685369039f32ff815ef60166686e0162f540f1520c", size = 207349, upload-time = "2026-03-05T19:55:35.911Z" }, + { url = "https://files.pythonhosted.org/packages/00/15/b92b4e1d88d02c6eff9733c9eea21846ab435cc4d813d84ccc5d335955df/graphql_core-3.2.11-py3-none-any.whl", hash = "sha256:0b3e35ff41e9adba53021ab0cef475eb18f57c7f53f0f2ca55567fbf3c537ea0", size = 214879, upload-time = "2026-06-05T13:45:21.245Z" }, ] [[package]] @@ -402,17 +402,17 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version < '3.11'" }, - { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, - { name = "jedi", marker = "python_full_version < '3.11'" }, - { name = "matplotlib-inline", marker = "python_full_version < '3.11'" }, - { name = "pexpect", marker = "python_full_version < '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version < '3.11'" }, - { name = "pygments", marker = "python_full_version < '3.11'" }, - { name = "stack-data", marker = "python_full_version < '3.11'" }, - { name = "traitlets", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "exceptiongroup" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/40/18/f8598d287006885e7136451fdea0755af4ebcbfe342836f24deefaed1164/ipython-8.39.0.tar.gz", hash = "sha256:4110ae96012c379b8b6db898a07e186c40a2a1ef5d57a7fa83166047d9da7624", size = 5513971, upload-time = "2026-03-27T10:02:13.94Z" } wheels = [ @@ -427,17 +427,17 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "colorama", marker = "python_full_version == '3.11.*' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version == '3.11.*'" }, - { name = "ipython-pygments-lexers", marker = "python_full_version == '3.11.*'" }, - { name = "jedi", marker = "python_full_version == '3.11.*'" }, - { name = "matplotlib-inline", marker = "python_full_version == '3.11.*'" }, - { name = "pexpect", marker = "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version == '3.11.*'" }, - { name = "pygments", marker = "python_full_version == '3.11.*'" }, - { name = "stack-data", marker = "python_full_version == '3.11.*'" }, - { name = "traitlets", marker = "python_full_version == '3.11.*'" }, - { name = "typing-extensions", marker = "python_full_version == '3.11.*'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c5/25/daae0e764047b0a2480c7bbb25d48f4f509b5818636562eeac145d06dfee/ipython-9.10.1.tar.gz", hash = "sha256:e170e9b2a44312484415bdb750492699bf329233b03f2557a9692cce6466ada4", size = 4426663, upload-time = "2026-03-27T09:53:26.244Z" } wheels = [ @@ -452,16 +452,16 @@ resolution-markers = [ "python_full_version >= '3.12'", ] dependencies = [ - { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version >= '3.12'" }, - { name = "ipython-pygments-lexers", marker = "python_full_version >= '3.12'" }, - { name = "jedi", marker = "python_full_version >= '3.12'" }, - { name = "matplotlib-inline", marker = "python_full_version >= '3.12'" }, - { name = "pexpect", marker = "python_full_version >= '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version >= '3.12'" }, - { name = "pygments", marker = "python_full_version >= '3.12'" }, - { name = "stack-data", marker = "python_full_version >= '3.12'" }, - { name = "traitlets", marker = "python_full_version >= '3.12'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3a/73/7114f80a8f9cabdb13c27732dce24af945b2923dcab80723602f7c8bc2d8/ipython-9.12.0.tar.gz", hash = "sha256:01daa83f504b693ba523b5a407246cabde4eb4513285a3c6acaff11a66735ee4", size = 4428879, upload-time = "2026-03-27T09:42:45.312Z" } wheels = [ @@ -473,7 +473,7 @@ name = "ipython-pygments-lexers" version = "1.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pygments", marker = "python_full_version >= '3.11'" }, + { name = "pygments" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } wheels = [