For serializing models, just try dumping them with model_dump_json wh…

…ich calls different logic overall than to_jsonable_python for some reason I don't understand

qcelemental/models/basemodels.py

@@ -15,6 +15,10 @@ def _repr(self) -> str:
  return f'{self.__repr_name__()}({self.__repr_str__(", ")})'
 
 
+# Encoders, to be deprecated at some point
+ndarray_encoder = {np.ndarray: lambda v: v.flatten().tolist()}
+
+
 class ExtendedConfigDict(ConfigDict, total=False):
  serialize_default_excludes: Set
  """Add items to exclude from serialization"""

qcelemental/models/common_models.py

@@ -11,10 +11,6 @@
  ReprArgs = Sequence[Tuple[Optional[str], Any]]
 
 
-# Encoders, to be deprecated
-ndarray_encoder = {np.ndarray: lambda v: v.flatten().tolist()}
-
-
 def provenance_json_schema_extra(schema, model):
  schema["$schema"] = qcschema_draft
 

qcelemental/tests/test_utils.py

@@ -308,17 +308,83 @@ def test_serialization(obj, encoding):
  assert compare_recursive(obj, new_obj)
 
 
-@pytest.fixture(scope="function")
-def a_model():
- class M(qcel.models.ProtoModel):
- s: str = "1234"
- n: int = 12
- f: float = 4.0
- np: qcel.models.types.Array[float] = np.array([[1, 2, 3], [4, 5, 6]])
+@pytest.fixture
+def atomic_result():
+ """Mock AtomicResult output which can be tested against for complex serialization methods"""
+
+ data = {
+ "id": None,
+ "schema_name": "qcschema_output",
+ "schema_version": 1,
+ "molecule": {
+ "schema_name": "qcschema_molecule",
+ "schema_version": 2,
+ "validated": True,
+ "symbols": np.array(["O", "H", "H"], dtype="<U1"),
+ "geometry": np.array(
+ [
+ [0.0000000000000000, 0.0000000000000000, -0.1242978140796278],
+ [0.0000000000000000, -1.4344192748456206, 0.9863482549166890],
+ [0.0000000000000000, 1.4344192748456206, 0.9863482549166890],
+ ]
+ ),
+ "name": "h2o",
+ "molecular_charge": 0.0,
+ "molecular_multiplicity": 1,
+ "masses": np.array([15.9949146195700003, 1.0078250322300000, 1.0078250322300000]),
+ "real": np.array([True, True, True]),
+ "atom_labels": np.array(["", "", ""], dtype="<U1"),
+ "atomic_numbers": np.array([8, 1, 1], dtype=np.int16),
+ "mass_numbers": np.array([16, 1, 1], dtype=np.int16),
+ "fragments": [np.array([0, 1, 2], dtype=np.int32)],
+ "fragment_charges": [0.0],
+ "fragment_multiplicities": [1],
+ "fix_com": True,
+ "fix_orientation": True,
+ "provenance": {
+ "creator": "QCElemental",
+ "version": "0.29.0.dev1",
+ "routine": "qcelemental.molparse.from_string",
+ },
+ "extras": {},
+ },
+ "driver": "gradient",
+ "model": {"method": "unknown", "basis": "unknown"},
+ "keywords": {},
+ "protocols": {},
+ "extras": {
+ "qcvars": {
+ "NUCLEAR REPULSION ENERGY": 9.168193296424349,
+ "CURRENT ENERGY": -76.02663273512756,
+ "CURRENT GRADIENT": np.array(
+ [
+ [-0.0000000000000000, 0.0000000000000000, -0.0176416299024253],
+ [0.0000000000000000, -0.0124384148528182, 0.0088208149511995],
+ [-0.0000000000000000, 0.0124384148528182, 0.0088208149511995],
+ ]
+ ),
+ }
+ },
+ "provenance": {"creator": "User", "version": "0.1", "routine": ""},
+ "properties": {"nuclear_repulsion_energy": 9.168193296424349, "return_energy": -76.02663273512756},
+ "wavefunction": None,
+ "return_result": np.array(
+ [
+ [-0.0000000000000000, 0.0000000000000000, -0.0176416299024253],
+ [0.0000000000000000, -0.0124384148528182, 0.0088208149511995],
+ [-0.0000000000000000, 0.0124384148528182, 0.0088208149511995],
+ ]
+ ),
+ "stdout": "User provided energy, gradient, or hessian is returned",
+ "stderr": None,
+ "native_files": {},
+ "success": True,
+ "error": None,
+ }
 
- yield M()
+ yield qcel.models.results.AtomicResult(**data)
 
 
-def test_json_dumps(a_model):
- ret = qcel.util.json_dumps(a_model)
+def test_json_dumps(atomic_result):
+ ret = qcel.util.json_dumps(atomic_result)
  assert isinstance(ret, str)

qcelemental/util/serialization.py

@@ -2,7 +2,8 @@
 from typing import Any, Union
 
 import numpy as np
-from pydantic_core import to_jsonable_python
+import pydantic
+from pydantic_core import PydanticSerializationError, to_jsonable_python
 
 from .importing import which_import
 
@@ -192,11 +193,20 @@ def jsonext_loads(data: Union[str, bytes]) -> Any:
 
 class JSONArrayEncoder(json.JSONEncoder):
  def default(self, obj: Any) -> Any:
+ # See if pydantic can do this on its own.
+ # Note: This calls DIFFERENT logic on BaseModels than BaseModel.model_dump_json, for somoe reason
  try:
  return to_jsonable_python(obj)
  except ValueError:
  pass
 
+ # See if pydantic model can be just serialized if the above couldn't be dumped
+ if isinstance(obj, pydantic.BaseModel):
+ try:
+ return obj.model_dump_json()
+ except PydanticSerializationError:
+ pass
+
  if isinstance(obj, np.ndarray):
  if obj.shape:
  return obj.ravel().tolist()