diff --git a/.github/workflows/test_import.yml b/.github/workflows/test_import.yml index ae4d6efb..0677d330 100644 --- a/.github/workflows/test_import.yml +++ b/.github/workflows/test_import.yml @@ -15,4 +15,21 @@ jobs: architecture: 'x64' - run: python -m pip install uv - run: python -m uv pip install --system . - - run: python -c 'import dpdata' + - name: Test core import and YAML serialization + run: | + python - <<'PY' + import tempfile + from pathlib import Path + + import dpdata + + system = dpdata.System( + "tests/poscars/POSCAR.h2o.md", fmt="vasp/poscar" + ) + with tempfile.TemporaryDirectory() as tmpdir: + filename = Path(tmpdir) / "system.yaml" + system.dump(filename) + loaded = dpdata.System.load(filename) + assert loaded.formula == system.formula + assert loaded.get_nframes() == system.get_nframes() + PY diff --git a/AGENTS.md b/AGENTS.md index 19d633b9..d439826c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -9,7 +9,7 @@ Always reference these instructions first and fallback to search or bash command - **Bootstrap and install the repository:** - `cd /home/runner/work/dpdata/dpdata` (or wherever the repo is cloned) - - `uv pip install -e .` -- installs dpdata in development mode with core dependencies (numpy, scipy, h5py, monty, wcmatch) + - `uv pip install -e .` -- installs dpdata in development mode with core dependencies (numpy, scipy, h5py, PyYAML, wcmatch, lmdb, msgpack-numpy) - Test installation: `dpdata --version` -- should show version like "dpdata v0.1.dev2+..." - **Run tests:** @@ -93,7 +93,7 @@ The following are outputs from frequently run commands. Reference them instead o ### Key dependencies -- Core: numpy>=1.14.3, scipy, h5py, monty, wcmatch +- Core: numpy>=1.14.3, scipy, h5py, PyYAML, wcmatch, lmdb, msgpack-numpy - Optional: ase (ASE integration), parmed (AMBER), pymatgen (Materials Project), rdkit (molecular analysis) - Testing: unittest (built-in), coverage - Linting: ruff @@ -132,7 +132,7 @@ The following are outputs from frequently run commands. Reference them instead o - **Installation timeouts:** Network timeouts during `uv pip install` are common. If this occurs, try: - - Individual package installation: `uv pip install numpy scipy h5py monty wcmatch` + - Individual package installation: `uv pip install numpy scipy h5py PyYAML wcmatch lmdb msgpack-numpy` - Use `--timeout` option: `uv pip install --timeout 300 -e .` - Verify existing installation works: `dpdata --version` should work even if reinstall fails diff --git a/docs/conf.py b/docs/conf.py index 263cb550..ca38237e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -202,7 +202,6 @@ def setup(app): "numpy": ("https://docs.scipy.org/doc/numpy/", None), "python": ("https://docs.python.org/", None), "ase": ("https://wiki.fysik.dtu.dk/ase/", None), - "monty": ("https://guide.materialsvirtuallab.org/monty/", None), "h5py": ("https://docs.h5py.org/en/stable/", None), } diff --git a/docs/environment.yml b/docs/environment.yml index 2167efdd..5260c759 100644 --- a/docs/environment.yml +++ b/docs/environment.yml @@ -6,7 +6,6 @@ dependencies: - xeus-python - numpy - scipy - - monty - wcmatch - pip: - .. diff --git a/dpdata/serialization.py b/dpdata/serialization.py new file mode 100644 index 00000000..66afc30c --- /dev/null +++ b/dpdata/serialization.py @@ -0,0 +1,243 @@ +from __future__ import annotations + +import bz2 +import datetime +import gzip +import importlib +import json +import lzma +from enum import Enum +from pathlib import Path +from typing import Any, BinaryIO, TextIO, cast +from uuid import UUID + +import numpy as np + + +def _detect_format(filename: str | Path, fmt: str | None = None) -> str: + if fmt is not None: + return fmt + suffixes = [suffix.lower() for suffix in Path(filename).suffixes] + if suffixes and suffixes[-1] in {".gz", ".z", ".bz2", ".xz", ".lzma"}: + suffixes.pop() + suffix = suffixes[-1] if suffixes else "" + if suffix == ".mpk": + return "mpk" + if suffix in {".yaml", ".yml"}: + return "yaml" + return "json" + + +def _open_text(filename: str | Path, mode: str) -> TextIO: + path = str(filename) + lower_path = path.lower() + if lower_path.endswith((".gz", ".z")): + return cast("TextIO", gzip.open(path, mode, encoding="utf-8")) + if lower_path.endswith(".bz2"): + return cast("TextIO", bz2.open(path, mode, encoding="utf-8")) + if lower_path.endswith((".xz", ".lzma")): + return cast("TextIO", lzma.open(path, mode, encoding="utf-8")) + return cast("TextIO", open(path, mode, encoding="utf-8")) + + +def _open_binary(filename: str | Path, mode: str) -> BinaryIO: + path = str(filename) + lower_path = path.lower() + if lower_path.endswith((".gz", ".z")): + return cast("BinaryIO", gzip.open(path, mode)) + if lower_path.endswith(".bz2"): + return cast("BinaryIO", bz2.open(path, mode)) + if lower_path.endswith((".xz", ".lzma")): + return cast("BinaryIO", lzma.open(path, mode)) + return cast("BinaryIO", open(path, mode)) + + +def _yaml_dump(obj: Any, fp, *args: Any, **kwargs: Any) -> None: + try: + yaml = importlib.import_module("yaml") + + if "sort_keys" not in kwargs: + kwargs["sort_keys"] = False + getattr(yaml, "safe_dump")(obj, fp, *args, **kwargs) + except ModuleNotFoundError: + try: + ruamel_yaml = importlib.import_module("ruamel.yaml") + except ModuleNotFoundError as e: + raise RuntimeError( + "Dumping YAML files requires PyYAML or ruamel.yaml." + ) from e + yaml = getattr(ruamel_yaml, "YAML")() + if "indent" in kwargs: + indent = kwargs.pop("indent") + yaml.indent(mapping=indent, sequence=indent, offset=2) + yaml.dump(obj, fp, *args, **kwargs) + + +def _yaml_load(fp, *args: Any, **kwargs: Any) -> Any: + try: + yaml = importlib.import_module("yaml") + + return getattr(yaml, "safe_load")(fp, *args, **kwargs) + except ModuleNotFoundError: + try: + ruamel_yaml = importlib.import_module("ruamel.yaml") + except ModuleNotFoundError as e: + raise RuntimeError( + "Loading YAML files requires PyYAML or ruamel.yaml." + ) from e + yaml = getattr(ruamel_yaml, "YAML")(typ="safe") + return yaml.load(fp, *args, **kwargs) + + +def _encode_ndarray(obj: np.ndarray) -> dict[str, Any]: + if str(obj.dtype).startswith("complex"): + data = [obj.real.tolist(), obj.imag.tolist()] + else: + data = obj.tolist() + return { + "@module": "numpy", + "@class": "array", + "dtype": str(obj.dtype), + "data": data, + } + + +def to_serializable(obj: Any) -> Any: + """Convert common dpdata objects to monty-compatible plain data.""" + if isinstance(obj, dict): + return {to_serializable(k): to_serializable(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [to_serializable(v) for v in obj] + if isinstance(obj, np.ndarray): + return _encode_ndarray(obj) + if isinstance(obj, np.generic): + return obj.item() + if isinstance(obj, datetime.datetime): + return { + "@module": "datetime", + "@class": "datetime", + "string": obj.isoformat(), + } + if isinstance(obj, UUID): + return {"@module": "uuid", "@class": "UUID", "string": str(obj)} + if isinstance(obj, Path): + return {"@module": "pathlib", "@class": "Path", "string": str(obj)} + if isinstance(obj, Enum): + return { + "@module": obj.__class__.__module__, + "@class": obj.__class__.__name__, + "value": to_serializable(obj.value), + } + if hasattr(obj, "as_dict"): + data = obj.as_dict() + if "@module" not in data: + data["@module"] = obj.__class__.__module__ + if "@class" not in data: + data["@class"] = obj.__class__.__name__ + return to_serializable(data) + return obj + + +def _decode_ndarray(data: dict[str, Any]) -> np.ndarray: + dtype = data["dtype"] + if dtype.startswith("complex"): + real, imag = data["data"] + return np.array(real, dtype=dtype) + np.array(imag, dtype=dtype) * 1j + return np.array(data["data"], dtype=dtype) + + +def process_decoded(obj: Any) -> Any: + """Decode monty-style dictionaries used by existing dpdata JSON files.""" + if isinstance(obj, dict): + if "@module" in obj and "@class" in obj: + module_name = obj["@module"] + class_name = obj["@class"] + if module_name == "numpy" and class_name == "array": + return _decode_ndarray(obj) + if module_name == "datetime" and class_name == "datetime": + try: + return datetime.datetime.fromisoformat(obj["string"]) + except ValueError: + value = obj["string"].split("+")[0] + try: + return datetime.datetime.strptime(value, "%Y-%m-%d %H:%M:%S.%f") + except ValueError: + return datetime.datetime.strptime(value, "%Y-%m-%d %H:%M:%S") + if module_name == "uuid" and class_name == "UUID": + return UUID(obj["string"]) + if module_name == "pathlib" and class_name == "Path": + return Path(obj["string"]) + try: + module = importlib.import_module(module_name) + cls = getattr(module, class_name) + except (AttributeError, ImportError, ModuleNotFoundError): + cls = None + if cls is not None: + # Decode fields before construction so custom objects receive the + # same nested Python values as top-level serialized objects. + data = { + k: process_decoded(v) + for k, v in obj.items() + if not k.startswith("@") + } + if hasattr(cls, "from_dict"): + return cls.from_dict(data) + if isinstance(cls, type) and issubclass(cls, Enum): + return cls(data["value"]) + return {process_decoded(k): process_decoded(v) for k, v in obj.items()} + if isinstance(obj, list): + return [process_decoded(v) for v in obj] + return obj + + +def dumpfn( + obj: Any, + filename: str | Path, + *args: Any, + fmt: str | None = None, + **kwargs: Any, +) -> None: + """Dump an object to JSON, YAML, or msgpack without requiring monty.""" + fmt = _detect_format(filename, fmt) + obj = to_serializable(obj) + if fmt == "json": + with _open_text(filename, "wt") as fp: + json.dump(obj, fp, *args, **kwargs) + return + if fmt == "yaml": + with _open_text(filename, "wt") as fp: + _yaml_dump(obj, fp, *args, **kwargs) + return + if fmt == "mpk": + try: + import msgpack + except ModuleNotFoundError as e: + raise RuntimeError("Dumping msgpack files requires msgpack.") from e + with _open_binary(filename, "wb") as fp: + msgpack.dump(obj, fp, *args, **kwargs) + return + raise TypeError(f"Invalid format: {fmt}") + + +def loadfn( + filename: str | Path, + *args: Any, + fmt: str | None = None, + **kwargs: Any, +) -> Any: + """Load JSON, YAML, or msgpack data and decode monty-style objects.""" + fmt = _detect_format(filename, fmt) + if fmt == "json": + with _open_text(filename, "rt") as fp: + return process_decoded(json.load(fp, *args, **kwargs)) + if fmt == "yaml": + with _open_text(filename, "rt") as fp: + return process_decoded(_yaml_load(fp, *args, **kwargs)) + if fmt == "mpk": + try: + import msgpack + except ModuleNotFoundError as e: + raise RuntimeError("Loading msgpack files requires msgpack.") from e + with _open_binary(filename, "rb") as fp: + return process_decoded(msgpack.load(fp, *args, **kwargs)) + raise TypeError(f"Invalid format: {fmt}") diff --git a/dpdata/system.py b/dpdata/system.py index 18b16aaf..ecc465c6 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -332,10 +332,13 @@ def __add__(self, others): return self.__class__.from_dict({"data": self_copy.data}) def dump(self, filename: str, indent: int = 4): - """Dump .json or .yaml file.""" - from monty.serialization import dumpfn + """Dump a JSON, YAML, or MessagePack file.""" + from dpdata.serialization import _detect_format, dumpfn - dumpfn(self.as_dict(), filename, indent=indent) + if _detect_format(filename) == "mpk": + dumpfn(self.as_dict(), filename) + else: + dumpfn(self.as_dict(), filename, indent=indent) def map_atom_types( self, type_map: dict[str, int] | list[str] | None = None @@ -379,20 +382,18 @@ def map_atom_types( @staticmethod def load(filename: str): - """Rebuild System obj. from .json or .yaml file.""" - from monty.serialization import loadfn + """Rebuild a System object from a JSON, YAML, or MessagePack file.""" + from dpdata.serialization import loadfn return loadfn(filename) @classmethod def from_dict(cls, data: dict): """Construct a System instance from a data dict.""" - from monty.serialization import MontyDecoder # type: ignore + from dpdata.serialization import process_decoded decoded = { - k: MontyDecoder().process_decoded(v) - for k, v in data.items() - if not k.startswith("@") + k: process_decoded(v) for k, v in data.items() if not k.startswith("@") } return cls(**decoded) diff --git a/pyproject.toml b/pyproject.toml index b94fce2e..521b6b72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,9 +20,9 @@ classifiers = [ ] dependencies = [ 'numpy>=1.14.3', - 'monty', 'scipy', 'h5py', + 'PyYAML', 'wcmatch', 'lmdb', 'msgpack-numpy', @@ -121,7 +121,6 @@ banned-module-level-imports = [ "deepmd", "h5py", "wcmatch", - "monty", "scipy", ] diff --git a/tests/test_json.py b/tests/test_json.py index 0b6f1b9d..78404edb 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -1,10 +1,38 @@ from __future__ import annotations +import datetime +import os +import tempfile import unittest +import numpy as np from comp_sys import CompLabeledSys, IsPBC from context import dpdata +from dpdata.serialization import ( + _detect_format, + dumpfn, + loadfn, + process_decoded, + to_serializable, +) + + +class NestedSerializable: + """Test helper that records values passed through ``from_dict``.""" + + def __init__(self, value): + self.value = value + + def as_dict(self): + """Expose nested values through the same serialization hook as dpdata objects.""" + return {"value": self.value} + + @classmethod + def from_dict(cls, data): + """Construct the helper from decoded serialized data.""" + return cls(data["value"]) + class TestJsonLoad(unittest.TestCase, CompLabeledSys, IsPBC): def setUp(self): @@ -26,5 +54,101 @@ def setUp(self): self.v_places = 4 +class TestJsonDumpLoad(unittest.TestCase, CompLabeledSys, IsPBC): + def setUp(self): + self.system_1 = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") + self.tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(self.tmpdir.cleanup) + self.filename = os.path.join(self.tmpdir.name, "h2o.md.json") + self.system_1.dump(self.filename) + self.system_2 = dpdata.LabeledSystem.load(self.filename) + self.places = 6 + self.e_places = 6 + self.f_places = 6 + self.v_places = 4 + + +class TestYamlDumpLoad(unittest.TestCase, CompLabeledSys, IsPBC): + def setUp(self): + self.system_1 = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") + self.tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(self.tmpdir.cleanup) + self.filename = os.path.join(self.tmpdir.name, "h2o.md.yaml") + self.system_1.dump(self.filename) + self.system_2 = dpdata.LabeledSystem.load(self.filename) + self.places = 6 + self.e_places = 6 + self.f_places = 6 + self.v_places = 4 + + +class TestMessagePackDumpLoad(unittest.TestCase, CompLabeledSys, IsPBC): + def setUp(self): + self.system_1 = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") + self.tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(self.tmpdir.cleanup) + self.filename = os.path.join(self.tmpdir.name, "h2o.md.mpk") + self.system_1.dump(self.filename) + self.system_2 = dpdata.LabeledSystem.load(self.filename) + self.places = 6 + self.e_places = 6 + self.f_places = 6 + self.v_places = 4 + + +class TestSerialization(unittest.TestCase): + def test_detect_format_uses_final_meaningful_suffix(self): + self.assertEqual(_detect_format("data.mpk.gz"), "mpk") + self.assertEqual(_detect_format("data.yaml.bz2"), "yaml") + self.assertEqual(_detect_format("data.yaml.json"), "json") + self.assertEqual(_detect_format("data.mpk.backup"), "json") + + def test_from_dict_receives_decoded_nested_values(self): + value = NestedSerializable( + { + "array": np.array([1, 2, 3]), + "datetime": datetime.datetime(2026, 7, 18, 12, 0), + } + ) + + decoded = process_decoded(to_serializable(value)) + + np.testing.assert_array_equal(decoded.value["array"], value.value["array"]) + self.assertEqual(decoded.value["datetime"], value.value["datetime"]) + + def test_timezone_aware_datetime_round_trip(self): + for value in ( + datetime.datetime( + 2026, + 6, + 19, + 12, + 0, + 0, + 123456, + tzinfo=datetime.timezone(datetime.timedelta(hours=8)), + ), + datetime.datetime( + 2026, + 6, + 19, + 12, + 0, + 0, + tzinfo=datetime.timezone(datetime.timedelta(hours=-5)), + ), + ): + serialized = to_serializable(value) + self.assertEqual(serialized["string"], value.isoformat()) + self.assertEqual(process_decoded(serialized), value) + + def test_yaml_dump_load(self): + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, "data.yaml") + value = {"numbers": [1, 2, 3]} + dumpfn(value, filename) + self.assertEqual(loadfn(filename), value) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_serialization.py b/tests/test_serialization.py new file mode 100644 index 00000000..45a897b3 --- /dev/null +++ b/tests/test_serialization.py @@ -0,0 +1,335 @@ +from __future__ import annotations + +import datetime +import importlib +import os +import shutil +import sys +import tempfile +import unittest +import uuid +from enum import Enum +from pathlib import Path +from unittest import mock + +import numpy as np +from context import dpdata # noqa: F401 + +from dpdata.serialization import ( + _detect_format, + _open_binary, + _open_text, + dumpfn, + loadfn, + process_decoded, + to_serializable, +) + + +class Color(Enum): + RED = "red" + BLUE = "blue" + + +class TestDetectFormat(unittest.TestCase): + def test_explicit_format_wins(self): + self.assertEqual(_detect_format("data.json", fmt="yaml"), "yaml") + + def test_suffix_dispatch(self): + for name, expected in ( + ("data.json", "json"), + ("data.mpk", "mpk"), + ("data.yaml", "yaml"), + ("data.yml", "yaml"), + ("data", "json"), + ("data.unknown", "json"), + ): + with self.subTest(name=name): + self.assertEqual(_detect_format(name), expected) + + def test_compression_suffix_is_stripped_first(self): + for name, expected in ( + ("data.yaml.gz", "yaml"), + ("data.mpk.bz2", "mpk"), + ("data.json.z", "json"), + ("data.yaml.xz", "yaml"), + ("data.mpk.lzma", "mpk"), + ("data.gz", "json"), + ): + with self.subTest(name=name): + self.assertEqual(_detect_format(name), expected) + + +class TestCompressedStreams(unittest.TestCase): + def setUp(self): + self.tmp_dir = tempfile.mkdtemp() + + def tearDown(self): + shutil.rmtree(self.tmp_dir, ignore_errors=True) + + def _path(self, name): + return os.path.join(self.tmp_dir, name) + + def test_text_roundtrip_through_each_codec(self): + for name in ( + "plain.txt", + "gzipped.gz", + "packed.z", + "squeezed.bz2", + "compressed.xz", + "legacy.lzma", + ): + with self.subTest(name=name): + path = self._path(name) + with _open_text(path, "wt") as fp: + fp.write("hello") + with _open_text(path, "rt") as fp: + self.assertEqual(fp.read(), "hello") + + def test_binary_roundtrip_through_each_codec(self): + for name in ( + "plain.bin", + "gzipped.gz", + "packed.z", + "squeezed.bz2", + "compressed.xz", + "legacy.lzma", + ): + with self.subTest(name=name): + path = self._path(name) + with _open_binary(path, "wb") as fp: + fp.write(b"\x00\x01") + with _open_binary(path, "rb") as fp: + self.assertEqual(fp.read(), b"\x00\x01") + + def test_compressed_files_are_not_plain_text(self): + path = self._path("compressed.json.gz") + dumpfn({"a": 1}, path) + with open(path, "rb") as fp: + self.assertEqual(fp.read(2), b"\x1f\x8b") + self.assertEqual(loadfn(path), {"a": 1}) + + def test_xz_files_are_compressed_and_round_trip(self): + path = self._path("compressed.json.xz") + dumpfn({"a": 1}, path) + with open(path, "rb") as fp: + self.assertEqual(fp.read(6), b"\xfd7zXZ\x00") + self.assertEqual(loadfn(path), {"a": 1}) + + +class TestEncoding(unittest.TestCase): + def test_complex_array_roundtrip(self): + arr = np.array([1 + 2j, 3 - 4j]) + encoded = to_serializable(arr) + self.assertEqual(encoded["dtype"], str(arr.dtype)) + self.assertEqual(len(encoded["data"]), 2) # real and imaginary parts + np.testing.assert_allclose(process_decoded(encoded), arr) + + def test_real_array_roundtrip(self): + arr = np.arange(6, dtype=np.float32).reshape(2, 3) + np.testing.assert_allclose(process_decoded(to_serializable(arr)), arr) + + def test_numpy_scalar_becomes_python_scalar(self): + value = to_serializable(np.float64(1.5)) + self.assertIsInstance(value, float) + self.assertNotIsInstance(value, np.generic) + + def test_uuid_roundtrip(self): + value = uuid.uuid4() + self.assertEqual(process_decoded(to_serializable(value)), value) + + def test_path_roundtrip(self): + value = Path("a") / "b.txt" + self.assertEqual(process_decoded(to_serializable(value)), value) + + def test_enum_roundtrip(self): + encoded = to_serializable(Color.BLUE) + self.assertEqual(encoded["@class"], "Color") + self.assertIs(process_decoded(encoded), Color.BLUE) + + def test_unserializable_object_passes_through(self): + marker = object() + self.assertIs(to_serializable(marker), marker) + + def test_dict_keys_are_converted_too(self): + encoded = to_serializable({np.int64(3): np.int64(4)}) + self.assertEqual(encoded, {3: 4}) + + +class TestDatetimeDecoding(unittest.TestCase): + def test_aware_datetime_roundtrip(self): + value = datetime.datetime(2026, 7, 27, 12, 30, tzinfo=datetime.timezone.utc) + self.assertEqual(process_decoded(to_serializable(value)), value) + + def test_negative_offset_roundtrip(self): + value = datetime.datetime( + 2026, 7, 27, 12, 30, tzinfo=datetime.timezone(-datetime.timedelta(hours=5)) + ) + self.assertEqual(process_decoded(to_serializable(value)), value) + + def test_non_iso_payload_with_microseconds(self): + # `fromisoformat` demands zero-padded fields; the strptime fallbacks + # keep hand-edited or legacy payloads readable. + decoded = process_decoded( + { + "@module": "datetime", + "@class": "datetime", + "string": "2020-1-2 3:4:5.678900", + } + ) + self.assertEqual(decoded, datetime.datetime(2020, 1, 2, 3, 4, 5, 678900)) + + def test_non_iso_payload_without_microseconds(self): + decoded = process_decoded( + { + "@module": "datetime", + "@class": "datetime", + "string": "2020-1-2 3:4:5", + } + ) + self.assertEqual(decoded, datetime.datetime(2020, 1, 2, 3, 4, 5)) + + def test_non_iso_payload_with_offset_drops_the_offset(self): + # The fallback splits on "+", so an offset it cannot parse is lost + # rather than raising. + decoded = process_decoded( + { + "@module": "datetime", + "@class": "datetime", + "string": "2020-1-2 3:4:5+00:00", + } + ) + self.assertEqual(decoded, datetime.datetime(2020, 1, 2, 3, 4, 5)) + + def test_iso_payload_with_space_separator(self): + decoded = process_decoded( + { + "@module": "datetime", + "@class": "datetime", + "string": "2020-01-02 03:04:05.678900", + } + ) + self.assertEqual(decoded, datetime.datetime(2020, 1, 2, 3, 4, 5, 678900)) + + +class TestUnknownMarkers(unittest.TestCase): + def test_unimportable_module_is_left_as_a_dict(self): + payload = {"@module": "no.such.module", "@class": "Thing", "value": 1} + self.assertEqual(process_decoded(payload), payload) + + def test_missing_class_is_left_as_a_dict(self): + payload = {"@module": "datetime", "@class": "NotAClass", "value": 1} + self.assertEqual(process_decoded(payload), payload) + + def test_class_without_from_dict_is_left_as_a_dict(self): + payload = {"@module": "builtins", "@class": "object", "value": 1} + self.assertEqual(process_decoded(payload), payload) + + def test_nested_markers_inside_lists_are_decoded(self): + arr = np.arange(3) + decoded = process_decoded([to_serializable(arr), {"k": to_serializable(arr)}]) + np.testing.assert_allclose(decoded[0], arr) + np.testing.assert_allclose(decoded[1]["k"], arr) + + +class TestFileFormats(unittest.TestCase): + def setUp(self): + self.tmp_dir = tempfile.mkdtemp() + self.payload = { + "array": np.arange(4, dtype=np.float64), + "when": datetime.datetime(2026, 7, 27, 8, 0), + "nested": {"list": [1, 2, 3]}, + } + + def tearDown(self): + shutil.rmtree(self.tmp_dir, ignore_errors=True) + + def _path(self, name): + return os.path.join(self.tmp_dir, name) + + def _assert_roundtrip(self, path): + dumpfn(self.payload, path) + loaded = loadfn(path) + np.testing.assert_allclose(loaded["array"], self.payload["array"]) + self.assertEqual(loaded["when"], self.payload["when"]) + self.assertEqual(loaded["nested"], self.payload["nested"]) + + def test_json_roundtrip(self): + self._assert_roundtrip(self._path("data.json")) + + def test_yaml_roundtrip(self): + self._assert_roundtrip(self._path("data.yaml")) + + def test_msgpack_roundtrip(self): + self._assert_roundtrip(self._path("data.mpk")) + + def test_compressed_msgpack_roundtrip(self): + self._assert_roundtrip(self._path("data.mpk.gz")) + + def test_invalid_format_on_dump(self): + with self.assertRaisesRegex(TypeError, "Invalid format: toml"): + dumpfn({"a": 1}, self._path("data.toml"), fmt="toml") + + def test_invalid_format_on_load(self): + path = self._path("data.json") + dumpfn({"a": 1}, path) + with self.assertRaisesRegex(TypeError, "Invalid format: toml"): + loadfn(path, fmt="toml") + + +def _import_without(*blocked): + """Return an ``import_module`` that pretends ``blocked`` is not installed.""" + real = importlib.import_module + + def fake(name, *args, **kwargs): + if name in blocked: + raise ModuleNotFoundError(f"No module named {name!r}") + return real(name, *args, **kwargs) + + return fake + + +class TestOptionalDependencies(unittest.TestCase): + def setUp(self): + self.tmp_dir = tempfile.mkdtemp() + + def tearDown(self): + shutil.rmtree(self.tmp_dir, ignore_errors=True) + + def _path(self, name): + return os.path.join(self.tmp_dir, name) + + def test_yaml_falls_back_to_ruamel(self): + path = self._path("data.yaml") + with mock.patch("importlib.import_module", _import_without("yaml")): + dumpfn({"a": [1, 2]}, path, indent=2) + with mock.patch("importlib.import_module", _import_without("yaml")): + loaded = loadfn(path) + self.assertEqual(loaded, {"a": [1, 2]}) + + def test_yaml_without_any_backend_explains_itself(self): + blocked = _import_without("yaml", "ruamel.yaml") + with mock.patch("importlib.import_module", blocked): + with self.assertRaisesRegex(RuntimeError, "requires PyYAML or ruamel"): + dumpfn({"a": 1}, self._path("data.yaml")) + + path = self._path("readable.yaml") + dumpfn({"a": 1}, path) + with mock.patch("importlib.import_module", blocked): + with self.assertRaisesRegex(RuntimeError, "requires PyYAML or ruamel"): + loadfn(path) + + def test_msgpack_absence_explains_itself(self): + path = self._path("data.mpk") + with mock.patch.dict(sys.modules, {"msgpack": None}): + with self.assertRaisesRegex(RuntimeError, "requires msgpack"): + dumpfn({"a": 1}, path) + + dumpfn({"a": 1}, path) + with mock.patch.dict(sys.modules, {"msgpack": None}): + with self.assertRaisesRegex(RuntimeError, "requires msgpack"): + loadfn(path) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_to_pymatgen_entry.py b/tests/test_to_pymatgen_entry.py index dfdeb468..512b86b8 100644 --- a/tests/test_to_pymatgen_entry.py +++ b/tests/test_to_pymatgen_entry.py @@ -4,7 +4,8 @@ import unittest from context import dpdata -from monty.serialization import loadfn # noqa: TID253 + +from dpdata.serialization import loadfn try: from pymatgen.entries.computed_entries import ComputedStructureEntry # noqa: F401