Migrate to pydantic v2 (#182)

Update dependencies to support pydantic v2.
Run bump-pydantic and use typing.Annotated.

Fix custom OPTIMADEUrl type.
Update everything to the pydantic v2 way, meaning use classmethods
__get_pydantic_core_schema__ and __get_pydantic_json_schema__,
drastically minimizing the class, but also making it a bit more opaque.

Use new pydantic v2 BaseModel method names.

Remove dependency on sessions containing Python objects - rely only on
pure JSON types in the session object.

Remove model2dict() utility function.
What it does (recursively dump all models as dicts) is the default
behaviour in pydantic v2 now, so there is no need for this
functionality.

Add scheme to OPTIMADEUrl
It was found to be necessary when used in certain core OTEAPI
strategies.
It may be relevant to have all the pydantic_core.Url properties readily
available.
Also, make the OPTIMADE accessUrl further variable by making the host a
changeable environment variable (OPTIMADE_HOST).

Add Python 3.11 support.
Cap supported Python versions at 3.11 due to DLite not supporting Python
3.12 and above.

Use oteapi-services version prior to session/parser upgrade.

Author: CasperWA

.github/utils/end2end_test.py

@@ -67,7 +67,7 @@ def main(oteapi_url: str) -> None:
 
     source = client.create_dataresource(
         accessService="OPTIMADE",
-        accessUrl=f"http://localhost:{os.getenv('OPTIMADE_PORT', '5000')}/",
+        accessUrl=OPTIMADE_URL,
         configuration=config,
     )
 
@@ -133,6 +133,7 @@ def main(oteapi_url: str) -> None:
     PORT = os.getenv("OTEAPI_PORT", "8080")
     OTEAPI_SERVICE_URL = f"http://localhost:{PORT}"
     OTEAPI_PREFIX = os.getenv("OTEAPI_prefix", "/api/v1")  # noqa: SIM112
+    OPTIMADE_URL = f"http://{os.getenv('OPTIMADE_HOST', 'localhost')}:{os.getenv('OPTIMADE_PORT', '5000')}/"
     if "OTEAPI_prefix" not in os.environ:
         # Set environment variables
         os.environ["OTEAPI_prefix"] = OTEAPI_PREFIX  # noqa: SIM112

.github/workflows/ci_tests.yml

@@ -40,7 +40,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10"]
+        python-version: ["3.9", "3.10", "3.11"]
         os:
         - ["ubuntu-latest", "linux"]
         - ["windows-latest", "windows"]
@@ -132,9 +132,11 @@ jobs:
         .github/utils/wait_for_it.sh localhost:${OTEAPI_PORT} -t 240
         sleep 5
       env:
-        # Pin to 1.20231108.329 until #163 has been closed.
-        # Related issue: https://github.com/SINTEF/oteapi-optimade/issues/187
-        DOCKER_OTEAPI_VERSION: '1.20231108.329'
+        # Use version 1.20240228.345 until
+        # https://github.com/SINTEF/oteapi-optimade/issues/213 has been resolved.
+        # See also
+        # https://github.com/EMMC-ASBL/oteapi-services/tree/8306d7212419764fb87e5cefdb5a869db9c68ef7?tab=readme-ov-file#open-translation-environment-ote-api
+        DOCKER_OTEAPI_VERSION: '1.20240228.345'
 
     - name: Run end-2-end tests
       run: python .github/utils/end2end_test.py

.pre-commit-config.yaml

@@ -54,7 +54,7 @@ repos:
       exclude: ^tests/.*$
       additional_dependencies:
         - "types-requests"
-        - "pydantic<2"
+        - "pydantic>=2,<3"
 
   - repo: https://github.com/SINTEF/ci-cd
     rev: v2.7.2

docs/api_reference/utils.md

@@ -8,6 +8,10 @@
 
 from __future__ import annotations
 
+import logging
+
 __version__ = "0.4.2"
 __author__ = "Casper Welzel Andersen"
 __author_email__ = "casper.w.andersen@sintef.no"
+
+logging.getLogger("oteapi_optimade").setLevel(logging.DEBUG)

oteapi_optimade/__init__.py

@@ -2,13 +2,21 @@
 
 from __future__ import annotations
 
+import importlib
 import logging
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import dlite
 from optimade.adapters import Structure
-from optimade.models import StructureResponseMany, StructureResponseOne, Success
+from optimade.models import (
+    Response as OPTIMADEResponse,
+)
+from optimade.models import (
+    StructureResponseMany,
+    StructureResponseOne,
+    Success,
+)
 from oteapi.models import SessionUpdate
 from oteapi_dlite.models import DLiteSessionUpdate
 from oteapi_dlite.utils import get_collection, update_collection
@@ -89,7 +97,14 @@ def get(
             context from services.
 
         """
-        session = OPTIMADEParseStrategy(self.parse_config).get(session)
+        generic_parse_config = self.parse_config.model_copy(
+            update={
+                "mediaType": self.parse_config.mediaType.lower().replace(
+                    "+dlite", "+json"
+                )
+            }
+        ).model_dump()
+        session = OPTIMADEParseStrategy(generic_parse_config).get(session)
 
         entities_path = Path(__file__).resolve().parent.resolve() / "entities"
 
@@ -111,108 +126,117 @@ def get(
             f"yaml://{entities_path}/OPTIMADEStructureSpecies.yaml"
         )
 
+        if not all(
+            _ in session for _ in ("optimade_response", "optimade_response_model")
+        ):
+            base_error_message = (
+                "Could not retrieve response from OPTIMADE parse strategy."
+            )
+            LOGGER.error(
+                "%s\n"
+                "optimade_response=%r\n"
+                "optimade_response_model=%r\n"
+                "session fields=%r",
+                base_error_message,
+                session.get("optimade_response"),
+                session.get("optimade_response_model"),
+                list(session.keys()),
+            )
+            raise OPTIMADEParseError(base_error_message)
+
+        optimade_response_model_module, optimade_response_model_name = session.get(
+            "optimade_response_model"
+        )
+        optimade_response_dict = session.get("optimade_response")
+
         error_message_supporting_only_structures = (
             "The DLite OPTIMADE Parser currently only supports structures entities."
         )
 
-        if self.parse_config.configuration.return_object:
-            # The response is given as a "proper" pydantic data model instance
-
-            if "optimade_response_object" not in session:
-                error_message = (
-                    "'optimade_response_object' was expected to be present in the "
-                    "session."
+        # Parse response using the provided model
+        try:
+            optimade_response_model: type[OPTIMADEResponse] = getattr(
+                importlib.import_module(optimade_response_model_module),
+                optimade_response_model_name,
+            )
+            optimade_response = optimade_response_model(**optimade_response_dict)
+        except (ImportError, AttributeError) as exc:
+            base_error_message = "Could not import the response model."
+            LOGGER.error(
+                "%s\n"
+                "ImportError: %s\n"
+                "optimade_response_model_module=%r\n"
+                "optimade_response_model_name=%r",
+                base_error_message,
+                exc,
+                optimade_response_model_module,
+                optimade_response_model_name,
+            )
+            raise OPTIMADEParseError(base_error_message) from exc
+        except ValidationError as exc:
+            base_error_message = "Could not validate the response model."
+            LOGGER.error(
+                "%s\n"
+                "ValidationError: %s\n"
+                "optimade_response_model_module=%r\n"
+                "optimade_response_model_name=%r",
+                base_error_message,
+                exc,
+                optimade_response_model_module,
+                optimade_response_model_name,
+            )
+            raise OPTIMADEParseError(base_error_message) from exc
+
+        # Currently, only "structures" entries are supported and handled
+        if isinstance(optimade_response, StructureResponseMany):
+            structures = [
+                (
+                    Structure(entry)
+                    if isinstance(entry, dict)
+                    else Structure(entry.model_dump())
                 )
-                raise ValueError(error_message)
-
-            # Currently, only "structures" entries are supported and handled
-            if isinstance(session.optimade_response_object, StructureResponseMany):
+                for entry in optimade_response.data
+            ]
+        elif isinstance(optimade_response, StructureResponseOne):
+            structures = [
+                (
+                    Structure(optimade_response.data)
+                    if isinstance(optimade_response.data, dict)
+                    else Structure(optimade_response.data.model_dump())
+                )
+            ]
+        elif isinstance(optimade_response, Success):
+            if isinstance(optimade_response.data, dict):
+                structures = [Structure(optimade_response.data)]
+            elif isinstance(optimade_response.data, BaseModel):
+                structures = [Structure(optimade_response.data.model_dump())]
+            elif isinstance(optimade_response.data, list):
                 structures = [
                     (
                         Structure(entry)
                         if isinstance(entry, dict)
-                        else Structure(entry.dict())
+                        else Structure(entry.model_dump())
                     )
-                    for entry in session.optimade_response_object.data
+                    for entry in optimade_response.data
                 ]
-            elif isinstance(session.optimade_response_object, StructureResponseOne):
-                structures = [
-                    (
-                        Structure(session.optimade_response_object.data)
-                        if isinstance(session.optimade_response_object.data, dict)
-                        else Structure(session.optimade_response_object.data.dict())
-                    )
-                ]
-            elif isinstance(session.optimade_response_object, Success):
-                if isinstance(session.optimade_response_object.data, dict):
-                    structures = [Structure(session.optimade_response_object.data)]
-                elif isinstance(session.optimade_response_object.data, BaseModel):
-                    structures = [
-                        Structure(session.optimade_response_object.data.dict())
-                    ]
-                elif isinstance(session.optimade_response_object.data, list):
-                    structures = [
-                        (
-                            Structure(entry)
-                            if isinstance(entry, dict)
-                            else Structure(entry.dict())
-                        )
-                        for entry in session.optimade_response_object.data
-                    ]
-                else:
-                    LOGGER.debug(
-                        "Could not determine what to do with `data`. Type %s.",
-                        type(session.optimade_response_object.data),
-                    )
-                    error_message = "Could not parse `data` entry in response."
-                    raise OPTIMADEParseError(error_message)
             else:
-                LOGGER.debug(
-                    "Got currently unsupported response type %s. Only structures are "
-                    "supported.",
-                    session.optimade_response_object.__class__.__name__,
+                LOGGER.error(
+                    "Could not determine what to do with `data`. Type %s.",
+                    type(optimade_response.data),
                 )
-                raise OPTIMADEParseError(error_message_supporting_only_structures)
+                error_message = "Could not parse `data` entry in response."
+                raise OPTIMADEParseError(error_message)
         else:
-            # The response is given as pure Python dictionary
-
-            if "optimade_response" not in session:
-                error_message = (
-                    "'optimade_response' was expected to be present in the session."
-                )
-                raise ValueError(error_message)
-
-            if not session.optimade_response or "data" not in session.optimade_response:
-                LOGGER.debug("Not a successful response - no 'data' entry found.")
-                return session
-
-            if isinstance(session.optimade_response["data"], list):
-                try:
-                    structures = [
-                        Structure(entry) for entry in session.optimade_response["data"]
-                    ]
-                except ValidationError as exc:
-                    LOGGER.debug(
-                        "Could not parse list of 'data' entries as structures."
-                    )
-                    raise OPTIMADEParseError(
-                        error_message_supporting_only_structures
-                    ) from exc
-            elif session.optimade_response is not None:
-                try:
-                    structures = [Structure(session.optimade_response["data"])]
-                except ValidationError as exc:
-                    LOGGER.debug("Could not parse single 'data' entry as a structure.")
-                    raise OPTIMADEParseError(
-                        error_message_supporting_only_structures
-                    ) from exc
-            else:
-                LOGGER.debug("Could not parse 'data' entries as structures.")  # type: ignore[unreachable]
-                raise OPTIMADEParseError(error_message_supporting_only_structures)
+            LOGGER.error(
+                "Got currently unsupported response type %s. Only structures are "
+                "supported.",
+                optimade_response_model_name,
+            )
+            raise OPTIMADEParseError(error_message_supporting_only_structures)
 
+        # DLite-fy OPTIMADE structures
         dlite_collection = get_collection(session)
 
-        # DLite-fy OPTIMADE structures
         for structure in structures:
             new_structure_attributes: dict[str, Any] = {}
 
@@ -226,7 +250,7 @@ def get(
                 for assembly in structure.attributes.assemblies:
                     # Ensure we're dealing with a normal Python dict
                     assembly_dict = (
-                        assembly.dict(exclude_none=True)
+                        assembly.model_dump(exclude_none=True)
                         if isinstance(assembly, BaseModel)
                         else assembly
                     )
@@ -252,7 +276,7 @@ def get(
                 for species_individual in structure.attributes.species:
                     # Ensure we're dealing with a normal Python dict
                     species_individual_dict = (
-                        species_individual.dict(exclude_none=True)
+                        species_individual.model_dump(exclude_none=True)
                         if isinstance(species_individual, BaseModel)
                         else species_individual
                     )
@@ -274,7 +298,7 @@ def get(
 
             # Attributes
             new_structure_attributes.update(
-                structure.attributes.dict(
+                structure.attributes.model_dump(
                     exclude={"species", "assemblies", "nelements", "nsites"}
                 )
             )