From e84be0130b46ac63e6e633931d04804f89fea665 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 20 Jun 2026 01:40:29 +0800 Subject: [PATCH 1/5] fix(serialization): remove monty dependency --- AGENTS.md | 6 +- docs/conf.py | 1 - docs/environment.yml | 1 - dpdata/serialization.py | 226 ++++++++++++++++++++++++++++++++ dpdata/system.py | 10 +- pyproject.toml | 2 - tests/test_json.py | 18 +++ tests/test_to_pymatgen_entry.py | 3 +- 8 files changed, 253 insertions(+), 14 deletions(-) create mode 100644 dpdata/serialization.py diff --git a/AGENTS.md b/AGENTS.md index 19d633b99..feb6116cd 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, wcmatch) - 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, wcmatch - 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 wcmatch` - 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 263cb5507..ca38237ed 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 89d2e5cad..14d5a7ee4 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 000000000..44b985671 --- /dev/null +++ b/dpdata/serialization.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import bz2 +import datetime +import gzip +import importlib +import json +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 + basename = Path(filename).name.lower() + if ".mpk" in basename: + return "mpk" + if ".yaml" in basename or ".yml" in basename: + 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")) + 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)) + 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": str(obj), + } + 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": + 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: + data = {k: 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(process_decoded(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 4150abc89..1ffc4f244 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -331,7 +331,7 @@ def __add__(self, others): def dump(self, filename: str, indent: int = 4): """Dump .json or .yaml file.""" - from monty.serialization import dumpfn + from dpdata.serialization import dumpfn dumpfn(self.as_dict(), filename, indent=indent) @@ -378,19 +378,17 @@ def map_atom_types( @staticmethod def load(filename: str): """Rebuild System obj. from .json or .yaml file.""" - from monty.serialization import loadfn + 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 b94fce2e9..b88e305a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,6 @@ classifiers = [ ] dependencies = [ 'numpy>=1.14.3', - 'monty', 'scipy', 'h5py', 'wcmatch', @@ -121,7 +120,6 @@ banned-module-level-imports = [ "deepmd", "h5py", "wcmatch", - "monty", "scipy", ] diff --git a/tests/test_json.py b/tests/test_json.py index 0b6f1b9dd..c98e21db2 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -1,5 +1,7 @@ from __future__ import annotations +import os +import tempfile import unittest from comp_sys import CompLabeledSys, IsPBC @@ -26,5 +28,21 @@ 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.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 + + def tearDown(self): + self.tmpdir.cleanup() + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_to_pymatgen_entry.py b/tests/test_to_pymatgen_entry.py index dfdeb4680..512b86b8b 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 From 1d9a382fccf1f6e2e43aed9a9a2f95d0a3112b5b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 19:38:50 +0800 Subject: [PATCH 2/5] fix(serialization): address review comments Preserve timezone-aware datetime values, retain YAML support through a declared backend, and align dependency documentation with project metadata. Add focused regression coverage for datetime and YAML round trips. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- AGENTS.md | 6 +++--- dpdata/serialization.py | 9 ++++++--- pyproject.toml | 1 + tests/test_json.py | 36 ++++++++++++++++++++++++++++++++++++ 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index feb6116cd..d439826c8 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, 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, 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 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/dpdata/serialization.py b/dpdata/serialization.py index 44b985671..8a5dc64d6 100644 --- a/dpdata/serialization.py +++ b/dpdata/serialization.py @@ -147,11 +147,14 @@ def process_decoded(obj: Any) -> Any: if module_name == "numpy" and class_name == "array": return _decode_ndarray(obj) if module_name == "datetime" and class_name == "datetime": - value = obj["string"].split("+")[0] try: - return datetime.datetime.strptime(value, "%Y-%m-%d %H:%M:%S.%f") + return datetime.datetime.fromisoformat(obj["string"]) except ValueError: - return datetime.datetime.strptime(value, "%Y-%m-%d %H:%M:%S") + 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": diff --git a/pyproject.toml b/pyproject.toml index b88e305a1..521b6b72b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ 'numpy>=1.14.3', 'scipy', 'h5py', + 'PyYAML', 'wcmatch', 'lmdb', 'msgpack-numpy', diff --git a/tests/test_json.py b/tests/test_json.py index c98e21db2..af02f5c69 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -1,5 +1,6 @@ from __future__ import annotations +import datetime import os import tempfile import unittest @@ -7,6 +8,8 @@ from comp_sys import CompLabeledSys, IsPBC from context import dpdata +from dpdata.serialization import dumpfn, loadfn, process_decoded, to_serializable + class TestJsonLoad(unittest.TestCase, CompLabeledSys, IsPBC): def setUp(self): @@ -44,5 +47,38 @@ def tearDown(self): self.tmpdir.cleanup() +class TestSerialization(unittest.TestCase): + 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)), + ), + ): + self.assertEqual(process_decoded(to_serializable(value)), 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() From b1e32a5a5107b15f5b899cac8e93de99d07fc1a3 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 18 Jul 2026 15:34:56 +0800 Subject: [PATCH 3/5] fix(serialization): address nested decoding reviews Detect serialization formats from the final meaningful suffix and recursively decode values before passing them to custom from_dict constructors. Add regression tests for both review findings. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- dpdata/serialization.py | 19 ++++++++++++++----- tests/test_json.py | 40 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 53 insertions(+), 6 deletions(-) diff --git a/dpdata/serialization.py b/dpdata/serialization.py index 8a5dc64d6..b7af831c5 100644 --- a/dpdata/serialization.py +++ b/dpdata/serialization.py @@ -16,10 +16,13 @@ def _detect_format(filename: str | Path, fmt: str | None = None) -> str: if fmt is not None: return fmt - basename = Path(filename).name.lower() - if ".mpk" in basename: + suffixes = [suffix.lower() for suffix in Path(filename).suffixes] + if suffixes and suffixes[-1] in {".gz", ".z", ".bz2"}: + suffixes.pop() + suffix = suffixes[-1] if suffixes else "" + if suffix == ".mpk": return "mpk" - if ".yaml" in basename or ".yml" in basename: + if suffix in {".yaml", ".yml"}: return "yaml" return "json" @@ -165,11 +168,17 @@ def process_decoded(obj: Any) -> Any: except (AttributeError, ImportError, ModuleNotFoundError): cls = None if cls is not None: - data = {k: v for k, v in obj.items() if not k.startswith("@")} + # 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(process_decoded(data["value"])) + 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] diff --git a/tests/test_json.py b/tests/test_json.py index af02f5c69..7238539f6 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -5,10 +5,29 @@ import tempfile import unittest +import numpy as np from comp_sys import CompLabeledSys, IsPBC from context import dpdata -from dpdata.serialization import dumpfn, loadfn, process_decoded, to_serializable +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 + + @classmethod + def from_dict(cls, data): + """Construct the helper from decoded serialized data.""" + return cls(data["value"]) class TestJsonLoad(unittest.TestCase, CompLabeledSys, IsPBC): @@ -48,6 +67,25 @@ def tearDown(self): 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( From 278fb697cb6ccf6c9c5d050f1e126d5eefaedca6 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 21 Jul 2026 14:07:29 +0800 Subject: [PATCH 4/5] test(serialization): exercise nested as_dict encoding Make the nested serialization helper expose as_dict so the regression test genuinely traverses NumPy and datetime values through the encoder. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- tests/test_json.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_json.py b/tests/test_json.py index 7238539f6..462ddf321 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -24,6 +24,10 @@ class NestedSerializable: 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.""" From adbdaf50fdbfbe5cd7f8c69e68d4dfbdb373fcd5 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 21 Jul 2026 14:23:59 +0800 Subject: [PATCH 5/5] fix(serialization): address final review comments Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- dpdata/serialization.py | 2 +- dpdata/system.py | 4 ++-- tests/test_json.py | 8 ++++---- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/dpdata/serialization.py b/dpdata/serialization.py index b7af831c5..b2fc8a630 100644 --- a/dpdata/serialization.py +++ b/dpdata/serialization.py @@ -111,7 +111,7 @@ def to_serializable(obj: Any) -> Any: return { "@module": "datetime", "@class": "datetime", - "string": str(obj), + "string": obj.isoformat(), } if isinstance(obj, UUID): return {"@module": "uuid", "@class": "UUID", "string": str(obj)} diff --git a/dpdata/system.py b/dpdata/system.py index c9e64ffee..a24acdc4e 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -332,7 +332,7 @@ 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.""" + """Dump a JSON, YAML, or MessagePack file.""" from dpdata.serialization import dumpfn dumpfn(self.as_dict(), filename, indent=indent) @@ -379,7 +379,7 @@ def map_atom_types( @staticmethod def load(filename: str): - """Rebuild System obj. from .json or .yaml file.""" + """Rebuild a System object from a JSON, YAML, or MessagePack file.""" from dpdata.serialization import loadfn return loadfn(filename) diff --git a/tests/test_json.py b/tests/test_json.py index 462ddf321..b44879dc6 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -58,6 +58,7 @@ 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) @@ -66,9 +67,6 @@ def setUp(self): self.f_places = 6 self.v_places = 4 - def tearDown(self): - self.tmpdir.cleanup() - class TestSerialization(unittest.TestCase): def test_detect_format_uses_final_meaningful_suffix(self): @@ -112,7 +110,9 @@ def test_timezone_aware_datetime_round_trip(self): tzinfo=datetime.timezone(datetime.timedelta(hours=-5)), ), ): - self.assertEqual(process_decoded(to_serializable(value)), value) + 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: