Defer aiida imports where not used (#79)

ml-evs · eimrek · web-flow · commit b8fcc0f34696 · 2025-08-15T11:57:53.000+02:00
* Only import aiida into config when type checking

* Defer import in conversion step too

* Run tests for 3.11 and 3.12

* Restructure aiida-plugin into config and main, and separate the functionality out

* Skip AiiDA test when not installed

* Define an extra without aiida and use it to run tests in CI without aiida available

* tiny ci fix

* skip serve test if aiida is not installed

---------

Co-authored-by: Kristjan Eimre &lt;kristjaneimre@gmail.com&gt;
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -16,8 +16,8 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  tests:
-    name: Test the `optimade-maker` package
+  linting:
+    name: Lint the `optimade-maker` package
     runs-on: ubuntu-latest
 
     steps:
@@ -31,7 +31,7 @@ jobs:
       - name: Set up uv
         uses: astral-sh/setup-uv@v5
         with:
-          version: "0.6.x"
+          version: "0.8.x"
           enable-cache: true
 
       - name: Install latest compatible versions of immediate dependencies
@@ -42,12 +42,46 @@ jobs:
         run: |
           uv run pre-commit run --all-files --show-diff-on-failure
 
+  tests:
+    name: Test the `optimade-maker` package
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: "${{ matrix.python-version }}"
+
+      - name: Set up uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          version: "0.8.x"
+          enable-cache: true
+
+      - name: Install latest compatible versions of immediate dependencies
+        run: |
+          uv sync --locked --all-extras --dev
+
       - name: Run tests
         run: |
           uv run pytest -vv --cov=./src/optimade_maker --cov-report=xml --cov-report=term ./tests
 
+      - name: Reinstall environment without aiida
+        run: |
+          uv sync --locked --extra core --extra tests
+
+      - name: Rerun tests without aiida
+        run: |
+          uv run pytest -vv --cov=./src/optimade_maker --cov-report=xml --cov-report=term ./tests --ignore=./tests/aiida
+
       - name: Upload coverage
         uses: codecov/codecov-action@v3
+        if: matrix.python-version == '3.10'
         with:
           name: project
           file: ./coverage.xml
diff --git a/pyproject.toml b/pyproject.toml
@@ -44,6 +44,7 @@ ase = ["ase ~= 3.22"]
 pymatgen = ["pymatgen >= 2023.9"]
 pandas = ["pandas >= 1.5, < 3"]
 aiida = ["aiida-core >= 2.6.3"]
+core = ["optimade-maker[ase,pymatgen,pandas]"]
 ingest = ["optimade-maker[ase,pymatgen,pandas,aiida]"]
 tests = ["pytest~=8.3", "pytest-cov~=6.0"]
 dev = ["ruff", "pre-commit", "mypy"]
diff --git a/src/optimade_maker/aiida_plugin/__init__.py b/src/optimade_maker/aiida_plugin/__init__.py
@@ -2,106 +2,47 @@
 
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
-
-import aiida
-from aiida import orm
-from aiida.common.exceptions import (
-    CorruptStorage,
-    IncompatibleStorageSchema,
-    NotExistent,
-    ProfileConfigurationError,
-    UnreachableStorage,
-)
-from aiida.storage.sqlite_zip.backend import SqliteZipBackend
+from typing import TYPE_CHECKING, Any
+
 from optimade.adapters import Structure
 from optimade.models import DataType, EntryResource
-from pydantic import BaseModel, Field, model_validator
 
 if TYPE_CHECKING:
-    from .config import EntryConfig
-
-
-class AiidaEntryPath(BaseModel):
-    """Config to specify an AiiDA entry path."""
-
-    aiida_file: Optional[str] = Field(
-        None, description="AiiDA file that contains the structures."
-    )
-    aiida_profile: Optional[str] = Field(
-        None, description="AiiDA profile that contains the structures."
-    )
-    aiida_group: Optional[str] = Field(
-        None,
-        description="AiiDA group that contains the structures. 'None' assumes all StructureData nodes",
-    )
+    from aiida import orm
 
-    @model_validator(mode="before")
-    @classmethod
-    def check_file_or_profile(cls, values):
-        if isinstance(values, list):
-            # Skip validation for lists
-            return values
-        if not values.get("aiida_file") and not values.get("aiida_profile"):
-            raise ValueError("Either 'aiida_file' or 'aiida_profile' must be defined.")
-        if values.get("aiida_file") and values.get("aiida_profile"):
-            raise ValueError(
-                "Both 'aiida_file' and 'aiida_profile' cannot be defined at the same time."
-            )
-        return values
+    from optimade_maker.config import EntryConfig
 
+from optimade_maker.aiida_plugin.config import AiidaEntryPath
 
-class AiidaQueryItem(BaseModel):
-    """An item representing a step in an AiiDA query, which allows
-        * to project properties of the current node, or
-        * to move to a connected node in the AiiDA provenance graph.
-    In the case of querying for connected nodes, the usual AiiDA
-    QueryBuilder filters and edge_filters can be applied.
-    """
 
-    project: Optional[str] = Field(
-        None, description="The AiiDA attribute to project in the query."
-    )
-    incoming_node: Optional[str] = Field(
-        None, description="Query for an incoming node of the specified type."
-    )
-    outgoing_node: Optional[str] = Field(
-        None, description="Query for an outgoing node of the specified type."
-    )
-    filters: Optional[dict[Any, Any]] = Field(
-        None, description="filters passed to AiiDA QueryBuilder."
-    )
-    edge_filters: Optional[dict[Any, Any]] = Field(
-        None, description="edge_filters passed to AiiDA QueryBuilder."
-    )
+def _check_aiida_import():
+    """Check if AiiDA is installed and raise an ImportError if not."""
+    try:
+        import aiida
 
-    @model_validator(mode="before")
-    @classmethod
-    def check_required_fields(cls, values):
-        if not any(
-            values.get(field) for field in ["project", "incoming_node", "outgoing_node"]
-        ):
-            raise ValueError(
-                "One of 'project', 'incoming_node', or 'outgoing_node' must be defined."
-            )
-        if values.get("filters") or values.get("edge_filters"):
-            if not any(
-                values.get(field) for field in ["incoming_node", "outgoing_node"]
-            ):
-                raise ValueError(
-                    "'filters' and 'edge_filters' can only be defined for 'incoming_node' or 'outgoing_node'."
-                )
-        return values
+        return aiida
+    except ImportError:
+        raise ImportError(
+            "The AiiDA plugin requires the `aiida-core` package to be installed. "
+            "Please install it with `pip install aiida-core`."
+        )
 
 
 def query_for_aiida_structures(
     structure_group: str | None = None,
-) -> dict[str, orm.StructureData]:
+) -> dict[str, "orm.StructureData"]:
     """
     Query for all aiida structures in the specified group
     (or all structures if no group specified)
     """
+    _check_aiida_import()
+    from aiida import orm
+    from aiida.common.exceptions import (
+        NotExistent,
+    )
+
     qb = orm.QueryBuilder()
+
     if structure_group:
         # check that the AiiDA group exists
         try:
@@ -123,6 +64,8 @@ def query_for_aiida_properties(
     Query for structure properties based on the custom aiida query format
     specified in the yaml file.
     """
+    aiida = _check_aiida_import()
+    from aiida import orm
 
     # query for the structures
     qb = orm.QueryBuilder()
@@ -185,6 +128,14 @@ def query_for_aiida_properties(
 
 
 def get_aiida_profile_from_file(path: Path):
+    _check_aiida_import()
+    from aiida.common.exceptions import (
+        CorruptStorage,
+        IncompatibleStorageSchema,
+        UnreachableStorage,
+    )
+    from aiida.storage.sqlite_zip.backend import SqliteZipBackend
+
     if not path.exists():
         raise FileNotFoundError(f"File not found: {path}")
 
@@ -200,6 +151,8 @@ def get_aiida_profile_from_file(path: Path):
 def convert_aiida_structure_to_optimade(
     aiida_structure: orm.StructureData,
 ) -> dict:
+    _check_aiida_import()
+
     try:
         ase_structure = aiida_structure.get_ase()
         optimade_entry = Structure.ingest_from(ase_structure).entry.model_dump()
@@ -242,6 +195,9 @@ def construct_entries_from_aiida(
     entry_config: EntryConfig,
     provider_prefix: str,
 ) -> dict[str, dict]:
+    aiida = _check_aiida_import()
+    from aiida.common.exceptions import ProfileConfigurationError
+
     if not isinstance(entry_config.entry_paths, AiidaEntryPath):
         raise RuntimeError("entry_paths is not AiiDA-specific.")
     if entry_config.entry_type != "structures":
@@ -278,4 +234,5 @@ def construct_entries_from_aiida(
                 optimade_entries[uuid]["attributes"][prop_name] = (
                     _convert_property_type(prop_def.type, prop)
                 )
+
     return optimade_entries
diff --git a/src/optimade_maker/aiida_plugin/config.py b/src/optimade_maker/aiida_plugin/config.py
@@ -0,0 +1,75 @@
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class AiidaEntryPath(BaseModel):
+    """Config to specify an AiiDA entry path."""
+
+    aiida_file: Optional[str] = Field(
+        None, description="AiiDA file that contains the structures."
+    )
+    aiida_profile: Optional[str] = Field(
+        None, description="AiiDA profile that contains the structures."
+    )
+    aiida_group: Optional[str] = Field(
+        None,
+        description="AiiDA group that contains the structures. 'None' assumes all StructureData nodes",
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_file_or_profile(cls, values):
+        if isinstance(values, list):
+            # Skip validation for lists
+            return values
+        if not values.get("aiida_file") and not values.get("aiida_profile"):
+            raise ValueError("Either 'aiida_file' or 'aiida_profile' must be defined.")
+        if values.get("aiida_file") and values.get("aiida_profile"):
+            raise ValueError(
+                "Both 'aiida_file' and 'aiida_profile' cannot be defined at the same time."
+            )
+        return values
+
+
+class AiidaQueryItem(BaseModel):
+    """An item representing a step in an AiiDA query, which allows
+        * to project properties of the current node, or
+        * to move to a connected node in the AiiDA provenance graph.
+    In the case of querying for connected nodes, the usual AiiDA
+    QueryBuilder filters and edge_filters can be applied.
+    """
+
+    project: Optional[str] = Field(
+        None, description="The AiiDA attribute to project in the query."
+    )
+    incoming_node: Optional[str] = Field(
+        None, description="Query for an incoming node of the specified type."
+    )
+    outgoing_node: Optional[str] = Field(
+        None, description="Query for an outgoing node of the specified type."
+    )
+    filters: Optional[dict[Any, Any]] = Field(
+        None, description="filters passed to AiiDA QueryBuilder."
+    )
+    edge_filters: Optional[dict[Any, Any]] = Field(
+        None, description="edge_filters passed to AiiDA QueryBuilder."
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_required_fields(cls, values):
+        if not any(
+            values.get(field) for field in ["project", "incoming_node", "outgoing_node"]
+        ):
+            raise ValueError(
+                "One of 'project', 'incoming_node', or 'outgoing_node' must be defined."
+            )
+        if values.get("filters") or values.get("edge_filters"):
+            if not any(
+                values.get(field) for field in ["incoming_node", "outgoing_node"]
+            ):
+                raise ValueError(
+                    "'filters' and 'edge_filters' can only be defined for 'incoming_node' or 'outgoing_node'."
+                )
+        return values
diff --git a/src/optimade_maker/config.py b/src/optimade_maker/config.py
@@ -10,12 +10,12 @@
 __version__ = "0.1.0"
 
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Union
 
 import yaml
 from pydantic import BaseModel, Field
 
-from .aiida_plugin import AiidaEntryPath, AiidaQueryItem
+from .aiida_plugin.config import AiidaEntryPath, AiidaQueryItem
 
 
 class UnsupportedConfigVersion(RuntimeError): ...
@@ -87,7 +87,7 @@ class EntryConfig(BaseModel):
     entry_type: str = Field(
         description="The OPTIMADE entry type, e.g. `structures` or `references`."
     )
-    entry_paths: list[ParsedFiles] | AiidaEntryPath = Field(
+    entry_paths: Union[list[ParsedFiles], AiidaEntryPath] = Field(
         description=(
             "A list of paths patterns to parse, provided relative to the top-level of the archive entry, "
             "after any compressed locations have been decompressed. Supports Python glob syntax for wildcards."
diff --git a/src/optimade_maker/convert.py b/src/optimade_maker/convert.py
@@ -16,7 +16,7 @@
 from optimade.models import EntryInfoResource, EntryResource
 from optimade.server.schemas import ENTRY_INFO_SCHEMAS, retrieve_queryable_properties
 
-from .aiida_plugin import AiidaEntryPath, construct_entries_from_aiida
+from .aiida_plugin.config import AiidaEntryPath
 from .config import Config, EntryConfig, JSONLConfig, ParsedFiles, PropertyDefinition
 
 PROVIDER_PREFIX = os.environ.get("OPTIMAKE_PROVIDER_PREFIX", "optimake")
@@ -613,11 +613,23 @@ def construct_entries(
     provider_prefix: str,
     limit: int | None = None,
 ) -> dict[str, dict]:
-    if isinstance(entry_config.entry_paths, AiidaEntryPath):
+    try:
+        is_aiida_archive = isinstance(entry_config.entry_paths, AiidaEntryPath)
+
+    except Exception:
+        is_aiida_archive = False
+
+    if is_aiida_archive:
+        try:
+            from .aiida_plugin import construct_entries_from_aiida
+        except ImportError:
+            raise RuntimeError(
+                "The AiiDA plugin is not installed. Please install it with `pip install optimake[aiida]`."
+            )
+
         optimade_entries = construct_entries_from_aiida(
             archive_path, entry_config, provider_prefix
         )
-
     else:
         optimade_entries = construct_entries_from_files(
             archive_path, entry_config, provider_prefix, limit=limit
diff --git a/tests/test_convert.py b/tests/test_convert.py
diff --git a/tests/test_serve.py b/tests/test_serve.py
diff --git a/uv.lock b/uv.lock
