Lazy import black and isort for faster test startup (#2556)

* Refactor: Optimize black and isort import handling with lazy loading

* Refactor: Add lazy import for inflect and clean up black mode import handling

* Refactor: Add lazy imports for black and isort in conftest.py

* Refactor: Improve isort configuration handling in apply_isort method

Author: koxudaxi

src/datamodel_code_generator/__main__.py

@@ -15,7 +15,6 @@
 from urllib.parse import ParseResult, urlparse
 
 import argcomplete
-import black
 from pydantic import BaseModel
 from typing_extensions import TypeAlias
 
@@ -35,6 +34,7 @@
     Formatter,
     PythonVersion,
     PythonVersionMin,
+    _get_black,
     is_supported_in_black,
 )
 from datamodel_code_generator.model.pydantic_v2 import UnionMode  # noqa: TC001 # needed for pydantic
@@ -518,7 +518,7 @@ def main(args: Sequence[str] | None = None) -> Exit:  # noqa: PLR0911, PLR0912,
         print(  # noqa: T201
             f"Installed black doesn't support Python version {config.target_python_version.value}.\n"
             f"You have to install a newer black.\n"
-            f"Installed black version: {black.__version__}",
+            f"Installed black version: {_get_black().__version__}",
             file=sys.stderr,
         )
         return Exit.ERROR

src/datamodel_code_generator/format.py

@@ -8,23 +8,41 @@
 
 import subprocess  # noqa: S404
 from enum import Enum
-from functools import cached_property
+from functools import cached_property, lru_cache
 from importlib import import_module
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 from warnings import warn
 
-import black
-import isort
-
 from datamodel_code_generator.util import load_toml
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
-try:
-    import black.mode
-except ImportError:  # pragma: no cover
-    black.mode = None
+
+
+@lru_cache(maxsize=1)
+def _get_black() -> Any:
+    import black as _black  # noqa: PLC0415
+
+    return _black
+
+
+@lru_cache(maxsize=1)
+def _get_black_mode() -> Any:  # pragma: no cover
+    black = _get_black()
+    try:
+        import black.mode  # noqa: PLC0415
+    except ImportError:
+        return None
+    else:
+        return black.mode
+
+
+@lru_cache(maxsize=1)
+def _get_isort() -> Any:
+    import isort as _isort  # noqa: PLC0415
+
+    return _isort
 
 
 class DatetimeClassType(Enum):
@@ -91,16 +109,19 @@ def has_strenum(self) -> bool:
 PythonVersionMin = PythonVersion.PY_39
 
 
-BLACK_PYTHON_VERSION: dict[PythonVersion, black.TargetVersion] = {
-    v: getattr(black.TargetVersion, f"PY{v.name.split('_')[-1]}")
-    for v in PythonVersion
-    if hasattr(black.TargetVersion, f"PY{v.name.split('_')[-1]}")
-}
+@lru_cache(maxsize=1)
+def _get_black_python_version_map() -> dict[PythonVersion, Any]:
+    black = _get_black()
+    return {
+        v: getattr(black.TargetVersion, f"PY{v.name.split('_')[-1]}")
+        for v in PythonVersion
+        if hasattr(black.TargetVersion, f"PY{v.name.split('_')[-1]}")
+    }
 
 
 def is_supported_in_black(python_version: PythonVersion) -> bool:  # pragma: no cover
     """Check if a Python version is supported by the installed black version."""
-    return python_version in BLACK_PYTHON_VERSION
+    return python_version in _get_black_python_version_map()
 
 
 def black_find_project_root(sources: Sequence[Path]) -> Path:
@@ -153,6 +174,10 @@ def __init__(  # noqa: PLR0912, PLR0913, PLR0917
         else:
             config = {}
 
+        black = _get_black()
+        black_mode = _get_black_mode()
+        isort = _get_isort()
+
         black_kwargs: dict[str, Any] = {}
         if wrap_string_literal is not None:
             experimental_string_processing = wrap_string_literal
@@ -175,10 +200,10 @@ def __init__(  # noqa: PLR0912, PLR0913, PLR0917
             elif experimental_string_processing:
                 black_kwargs["preview"] = True
                 black_kwargs["unstable"] = config.get("unstable", False)
-                black_kwargs["enabled_features"] = {black.mode.Preview.string_processing}
+                black_kwargs["enabled_features"] = {black_mode.Preview.string_processing}
 
         self.black_mode = black.FileMode(
-            target_versions={BLACK_PYTHON_VERSION[python_version]},
+            target_versions={_get_black_python_version_map()[python_version]},
             line_length=config.get("line-length", black.DEFAULT_LINE_LENGTH),
             string_normalization=not skip_string_normalization or not config.get("skip-string-normalization", True),
             **black_kwargs,
@@ -246,6 +271,7 @@ def format_code(
 
     def apply_black(self, code: str) -> str:
         """Format code using black."""
+        black = _get_black()
         return black.format_str(
             code,
             mode=self.black_mode,
@@ -271,27 +297,16 @@ def apply_ruff_formatter(self, code: str) -> str:
         )
         return result.stdout.decode(self.encoding)
 
-    if TYPE_CHECKING:
-
-        def apply_isort(self, code: str) -> str:
-            """Sort imports using isort."""
-            ...
-
-    elif isort.__version__.startswith("4."):
-
-        def apply_isort(self, code: str) -> str:
-            """Sort imports using isort v4."""
+    def apply_isort(self, code: str) -> str:
+        """Sort imports using isort."""
+        isort = _get_isort()
+        if self.isort_config is None:  # pragma: no cover
             return isort.SortImports(
                 file_contents=code,
                 settings_path=self.settings_path,
                 **self.isort_config_kwargs,
             ).output
-
-    else:
-
-        def apply_isort(self, code: str) -> str:
-            """Sort imports using isort v5+."""
-            return isort.code(code, config=self.isort_config)
+        return isort.code(code, config=self.isort_config)
 
 
 class CustomCodeFormatter:

tests/conftest.py

@@ -241,4 +241,8 @@ def _preload_heavy_modules() -> None:
     This reduces per-test overhead when running with pytest-xdist,
     as each worker only pays the import cost once at session start.
     """
+    import black  # noqa: PLC0415, F401
+    import inflect  # noqa: PLC0415, F401
+    import isort  # noqa: PLC0415, F401
+
     import datamodel_code_generator  # noqa: PLC0415, F401