model_rebuild is surprisingly slow, can we make it lazyloaded

[ghost]

We found that it takes more than 8 seconds to import the generated Client because we have more than 1k graphql models so that model_rebuild is called more than 1k times, which is super slow

python -X importtime -c "from graphql import Client"

....
import time:      1321 |       1321 |       pydantic._internal._std_types_schema
import time:      4473 |       5794 |     graphql.prompts_prompts_by_pk
import time:       198 |       5991 |   graphql.client
import time:      8467 |       8467 |   graphql.enums
import time:   7874752 |    7874752 |   graphql.input_types
import time:       760 |    8139339 | graphql

Is it possible to make model_rebuild lazyloaded?

[ghost]

import time drops significantly (8s -> 1.5s) which is acceptable in our use case

python -X importtime -c "from graphql import Client"

import time:      1304 |       1304 |       pydantic._internal._std_types_schema
import time:      3899 |       5203 |     graphql.prompts_prompts_by_pk
import time:       228 |       5430 |   graphql.client
import time:      8138 |       8138 |   graphql.enums
import time:   1275577 |    1275577 |   graphql.input_types
import time:       836 |    1526001 | graphql

if i modify the generated input_types.py a bit

class LazyInitMeta(type):
    _initialized = False

    def __call__(cls, *args, **kwargs):
        if not cls._initialized:
            cls.model_rebuild() # lazy init here
            cls._initialized = True
        return super().__call__(*args, **kwargs)

class CombinedMeta(LazyInitMeta, BaseModel.__class__):
    pass

# thousands of models...
class BigintComparisonExp(BaseModel, metaclass=CombinedMeta):
    eq: Optional[Any] = Field(alias="_eq", default=None)
    gt: Optional[Any] = Field(alias="_gt", default=None)
    gte: Optional[Any] = Field(alias="_gte", default=None)
    in_: Optional[List[Any]] = Field(alias="_in", default=None)
    is_null: Optional[bool] = Field(alias="_isNull", default=None)
    lt: Optional[Any] = Field(alias="_lt", default=None)
    lte: Optional[Any] = Field(alias="_lte", default=None)
    neq: Optional[Any] = Field(alias="_neq", default=None)
    nin: Optional[List[Any]] = Field(alias="_nin", default=None)

[mat-sop]

Hey, I haven't looked in-depth into the problem yet, but for now, you can reasonably easily automate your solution by using plugin hooks, e.g.

Separate LazyInitMeta and CombinedMeta to meta.py:

from .base_model import BaseModel

class LazyInitMeta(type):
    _initialized = False

    def __call__(cls, *args, **kwargs):
        if not cls._initialized:
            cls.model_rebuild() # lazy init here
            cls._initialized = True
        return super().__call__(*args, **kwargs)

class CombinedMeta(LazyInitMeta, BaseModel.__class__):
    pass

Include it in generated package:

# pyproject.toml
...
files_to_include = [".../meta.py"]

Define LazyLoadingPlugin in lazy_loading_plugin.py:

import ast

from ariadne_codegen.plugins.base import Plugin
from graphql import GraphQLInputObjectType


class LazyLoadingPlugin(Plugin):
    def generate_inputs_module(self, module: ast.Module) -> ast.Module:
        # adding "from .meta import CombinedMeta" to input_types.py
        module.body.insert(
            0,
            ast.ImportFrom(
                module="meta", names=[ast.alias(name="CombinedMeta")], level=1
            ),
        )

        # removing model_rebuild calls
        def _is_not_model_rebuild_call(stmt: ast.stmt) -> bool:
            return not (
                isinstance(stmt, ast.Expr)
                and isinstance(stmt.value, ast.Call)
                and isinstance(stmt.value.func, ast.Attribute)
                and stmt.value.func.attr == "model_rebuild"
            )

        module.body = list(filter(_is_not_model_rebuild_call, module.body))

        return module

    def generate_input_class(
        self, class_def: ast.ClassDef, input_type: GraphQLInputObjectType
    ) -> ast.ClassDef:
        # adding "metaclass=CombinedMeta" to every input class
        class_def.keywords.append(
            ast.keyword(arg="metaclass", value=ast.Name(id="CombinedMeta"))
        )

        return class_def

Add created plugin to configuration:

# pyproject.toml
...
plugins = ["....lazy_loading_plugin.LazyLoadingPlugin"]

[ghost]

Sounds great, though I don't know if we can delay those calls, would that be causing any other issues?

[ghost]

plugins = ["....lazy_loading_plugin.LazyLoadingPlugin"]

btw the plugin seems incorrect and i cannot figure out the correct path, in a python 3.11.2 virtual env

❯ tree
.
└── upservices
    ├── __init__.py
    ├── codegen.toml
    ├── gql
    │   └── prompts.gql
    ├── lazy_loading_plugin.py
    └── meta.py

2 directories, 5 files

❯ ariadne-codegen --config upservices/codegen.toml
Traceback (most recent call last):
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/bin/ariadne-codegen", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/lib/python3.11/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/lib/python3.11/site-packages/click/core.py", line 1055, in main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/lib/python3.11/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/lib/python3.11/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/lib/python3.11/site-packages/ariadne_codegen/main.py", line 33, in main
    client(config_dict)
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/lib/python3.11/site-packages/ariadne_codegen/main.py", line 56, in client
    plugins_types=get_plugins_types(settings.plugins),
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/lib/python3.11/site-packages/ariadne_codegen/plugins/explorer.py", line 13, in get_plugins_types
    if is_module_str(plugin_str):
       ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/chiminghuang/.virtualenvs/pymono-jobs-syjf/lib/python3.11/site-packages/ariadne_codegen/plugins/explorer.py", line 22, in is_module_str
    spec = importlib.util.find_spec(plugin_str)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib.util>", line 90, in find_spec
  File "<frozen importlib.util>", line 32, in resolve_name
ImportError: no package specified for '....lazy_loading_plugin.LazyLoadingPlugin' (required for relative module names)

❯ cat upservices/codegen.toml
───────┬───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
       │ File: upservices/codegen.toml
───────┼───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
   1   │ [tool.ariadne-codegen]
   2   │ remote_schema_url = "http://127.0.0.1:3600/graphql"
   3   │ queries_path = "upservices/gql/"
   4   │ target_package_name = "graphql"
   5   │ target_package_path = "upservices"
   6   │ include_comments = false
   7   │ files_to_include = ["upservices/meta.py"]
   8   │ plugins = ["....lazy_loading_plugin.LazyLoadingPlugin"]
───────┴──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────

[mat-sop]

Sounds great, though I don't know if we can delay those calls, would that be causing any other issues?

We're adding model_rebuild calls (update_forward_refs in pydantic<2) for recursive models which might be generated, depending on a specific schema. I haven't yet looked much into this, so I'm not yet sure if your solution is valid.

I found this issue on pydantic's gh, which suggests to me that there could be no documented change of behavior between update_forward_refs and model_rebuild, and we need check if those calls are still needed.

btw the plugin seems incorrect and i cannot figure out the correct path, in a python 3.11.2 virtual env

❯ tree
.
└── upservices
    ├── __init__.py
    ├── codegen.toml
    ├── gql
    │   └── prompts.gql
    ├── lazy_loading_plugin.py
    └── meta.py

"....lazy_loading_plugin.LazyLoadingPlugin" was just an example, you need to place lazy_loading_plugin.py in a module which is possible to import from.

In other words, this command cannot raise ImportError: python -c "from {module_name}.lazy_loading_plugin import LazyLoadingPlugin", and then place "{module_name}.lazy_loading_plugin.LazyLoadingPlugin" into config.

[mat-sop]

Closing because in #241, I completely removed model_rebuild calls.