[pull] main from dftd4:main

python/dftd4/qcschema.py

@@ -89,14 +89,35 @@
 -0.0002667885779142513
 """
 
-from typing import Union
+import sys
+from typing import Union, overload
 
 import numpy as np
-import qcelemental as qcel
 
 from .interface import DampingParam, DispersionModel
 from .library import get_api_version
 
+if sys.version_info < (3, 14):
+    try:
+        import qcelemental.models.v1 as qcel_v1
+    except ModuleNotFoundError:
+        import qcelemental.models as qcel_v1
+else:
+    qcel_v1 = None
+
+try:
+    import qcelemental.models.v2 as qcel_v2
+except ModuleNotFoundError:
+    qcel_v2 = None
+
+
+if qcel_v1 is None and qcel_v2 is None:
+    raise ModuleNotFoundError(
+        "The qcelemental package is required for qcschema support. "
+        "Please install it with 'pip install qcelemental'."
+    )
+
+
 _supported_drivers = [
     "energy",
     "gradient",
@@ -110,40 +131,57 @@
 _clean_dashlevel = str.maketrans("", "", "()")
 
 
-def run_qcschema(
-    input_data: Union[
-        dict, qcel.models.AtomicInput, "qcel.models.v2.AtomicInput"
-    ],
-) -> Union[qcel.models.AtomicResult, "qcel.models.v2.AtomicResult"]:
-    """Perform disperson correction based on an atomic inputmodel"""
+if qcel_v1 is not None:
+
+    @overload
+    def run_qcschema(
+        input_data: Union[dict, "qcel_v1.AtomicInput"],
+    ) -> "qcel_v1.AtomicResult": ...
 
-    v2_available = hasattr(qcel.models, "v2")
 
-    if v2_available and isinstance(input_data, qcel.models.v2.AtomicInput):
+if qcel_v2 is not None:
+
+    @overload
+    def run_qcschema(
+        input_data: Union[dict, "qcel_v2.AtomicInput"],
+    ) -> "qcel_v2.AtomicResult": ...
+
+
+def run_qcschema(input_data):
+    """Perform disperson correction based on an atomic inputmodel"""
+
+    if qcel_v2 is not None and isinstance(input_data, qcel_v2.AtomicInput):
         atomic_input = input_data
-    elif isinstance(input_data, qcel.models.AtomicInput):
+    elif qcel_v1 is not None and isinstance(input_data, qcel_v1.AtomicInput):
         atomic_input = input_data
-    elif v2_available and input_data.get("specification"):
-        atomic_input = qcel.models.v2.AtomicInput(**input_data)
+    elif qcel_v2 is not None and input_data.get("specification"):
+        atomic_input = qcel_v2.AtomicInput(**input_data)
+    elif qcel_v1 is not None:
+        atomic_input = qcel_v1.AtomicInput(**input_data)
     else:
-        atomic_input = qcel.models.AtomicInput(**input_data)
-
-    if (schver := atomic_input.schema_version) == 1:
-        from qcelemental.models import AtomicResult, ComputeError
-
-        ret_data = atomic_input.dict()
-    elif schver == 2:
-        from qcelemental.models.v2 import (
-            AtomicResult,
-            ComputeError,
-            FailedOperation,
+        raise ValueError(
+            "Input data is not a valid QCSchema AtomicInput for either v1 or v2."
         )
 
+    schema_version = atomic_input.schema_version
+    if schema_version == 1:
+        ret_data = atomic_input.dict()
+        input_keywords = atomic_input.keywords
+        input_method = atomic_input.model.method
+        input_driver = atomic_input.driver
+    elif schema_version == 2:
         ret_data = {
             "input_data": atomic_input,
             "extras": {},
             "molecule": atomic_input.molecule,
         }
+        input_keywords = atomic_input.specification.keywords
+        input_method = atomic_input.specification.model.method
+        input_driver = atomic_input.specification.driver
+    else:
+        raise ValueError(
+            f"Unsupported QCSchema version: {schema_version}. Only v1 and v2 are supported."
+        )
 
     provenance = {
         "creator": "dftd4",
@@ -156,43 +194,38 @@
 
     # Since it is a level hint we a forgiving if it is not present,
     # we are much less forgiving if the wrong level is hinted here.
-    atin_keywords = (
-        atomic_input.keywords
-        if schver == 1
-        else atomic_input.specification.keywords
-    )
-    _level = atin_keywords.get("level_hint", "d4")
+    _level = input_keywords.get("level_hint", "d4")
     if _level.lower() not in _available_levels:
-        error = ComputeError(
+        error = dict(
             error_type="input error",
-            error_message=f"Level '{_level}' is invalid for this dispersion correction",
+            error_message="Level '{}' is invalid for this dispersion correction".format(
+                _level
+            ),
         )
-        if schver == 1:
+        if schema_version == 1:
             ret_data.update(
                 provenance=provenance,
                 success=success,
                 properties=properties,
                 return_result=return_result,
-                error=error,
+                error=qcel_v1.ComputeError(**error),
+            )
+            return qcel_v1.AtomicResult(**ret_data)
+        elif schema_version == 2:
+            return qcel_v2.FailedOperation(
+                input_data=atomic_input, error=qcel_v2.ComputeError(**error)
             )
-            return AtomicResult(**ret_data)
-        elif schver == 2:
-            return FailedOperation(input_data=atomic_input, error=error)
 
     # Check if the method is provided and strip the “dashlevel” from the method
-    _method = (
-        atomic_input.model.method.split("-")
-        if schver == 1
-        else atomic_input.specification.model.method.split("-")
-    )
+    _method = input_method.split("-")
     if _method[-1].lower().translate(_clean_dashlevel) == _level.lower():
         _method.pop()
     _method = "-".join(_method)
     if len(_method) == 0:
         _method = None
 
     # Obtain the parameters for the damping function
-    _input_param = atin_keywords.get("params_tweaks", {"method": _method})
+    _input_param = input_keywords.get("params_tweaks", {"method": _method})
     if _level.lower() == "d4s":
         _model_param = {
             key: _input_param.pop(key, default)
@@ -222,16 +255,12 @@
             **_model_param,
         )
 
-        driver = (
-            atomic_input.driver
-            if schver == 1
-            else atomic_input.specification.driver
-        )
+        driver = input_driver
         res = disp.get_dispersion(
             param=param,
             grad=driver == "gradient",
         )
-        if atin_keywords.get("property", False):
+        if input_keywords.get("property", False):
             res.update(**disp.get_properties())
         extras = {"dftd4": res}
 
@@ -245,7 +274,7 @@
 
         properties.update(return_energy=res.get("energy"))
 
-        if atin_keywords.get("pair_resolved", False):
+        if input_keywords.get("pair_resolved", False):
             res = disp.get_pairwise_dispersion(param=param)
             extras["dftd4"].update(res)
 
@@ -256,7 +285,7 @@
             return_result = fullgrad
         else:
             ret_data.update(
-                error=ComputeError(
+                error=dict(
                     error_type="input error",
                     error_message="Calculation succeeded but invalid driver request provided",
                 ),
@@ -266,7 +295,7 @@
 
     except (RuntimeError, TypeError) as e:
         ret_data.update(
-            error=ComputeError(error_type="input error", error_message=str(e)),
+            error=dict(error_type="input error", error_message=str(e)),
         )
 
     ret_data.update(
@@ -276,7 +305,11 @@
         return_result=return_result,
     )
 
-    if schver == 2 and "error" in ret_data:
-        return FailedOperation(input_data=atomic_input, error=ret_data["error"])
+    if schema_version == 1:
+        return qcel_v1.AtomicResult(**ret_data)
 
-    return AtomicResult(**ret_data)
+    if "error" in ret_data:
+        return qcel_v2.FailedOperation(
+            input_data=atomic_input, error=ret_data["error"]
+        )
+    return qcel_v2.AtomicResult(**ret_data)