improve PymatgenTest.serialize_with_pickle: use self.tmp_path, add type hints and better err msg if test_eq fails

clean up tests using serialize_with_pickle()

Author: janosh

pymatgen/util/testing.py

@@ -9,17 +9,19 @@
 
 import json
 import string
-import tempfile
 import unittest
 from pathlib import Path
-from typing import ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import pytest
 from monty.json import MontyDecoder, MSONable
 from monty.serialization import loadfn
 
 from pymatgen.core import SETTINGS, Structure
 
+if TYPE_CHECKING:
+    from git import Sequence
+
 MODULE_DIR = Path(__file__).absolute().parent
 
 TEST_FILES_DIR = Path(SETTINGS.get("PMG_TEST_FILES_DIR", MODULE_DIR / ".." / ".." / "tests" / "files"))
@@ -51,6 +53,7 @@ def get_structure(cls, name: str) -> Structure:
             Structure
         """
         struct = cls.TEST_STRUCTURES.get(name) or loadfn(f"{cls.STRUCTURES_DIR}/{name}.json")
+        cls.TEST_STRUCTURES[name] = struct
         return struct.copy()
 
     @staticmethod
@@ -59,7 +62,7 @@ def assert_str_content_equal(actual, expected):
         strip_whitespace = {ord(c): None for c in string.whitespace}
         return actual.translate(strip_whitespace) == expected.translate(strip_whitespace)
 
-    def serialize_with_pickle(self, objects, protocols=None, test_eq=True):
+    def serialize_with_pickle(self, objects: Any, protocols: Sequence[int] = None, test_eq: bool = True):
         """Test whether the object(s) can be serialized and deserialized with
         pickle. This method tries to serialize the objects with pickle and the
         protocols specified in input. Then it deserializes the pickle format
@@ -77,7 +80,7 @@ def serialize_with_pickle(self, objects, protocols=None, test_eq=True):
             Nested list with the objects deserialized with the specified
             protocols.
         """
-        # Use the python version so that we get the traceback in case of errors
+        # use pickle, not cPickle so that we get the traceback in case of errors
         import pickle
 
         # Build a list even when we receive a single object.
@@ -86,39 +89,38 @@ def serialize_with_pickle(self, objects, protocols=None, test_eq=True):
             got_single_object = True
             objects = [objects]
 
-        if protocols is None:
-            protocols = [pickle.HIGHEST_PROTOCOL]
+        protocols = protocols or [pickle.HIGHEST_PROTOCOL]
 
-        # This list will contains the object deserialized with the different
-        # protocols.
+        # This list will contain the objects deserialized with the different protocols.
         objects_by_protocol, errors = [], []
 
         for protocol in protocols:
             # Serialize and deserialize the object.
-            mode = "wb"
-            fd, tmpfile = tempfile.mkstemp(text="b" not in mode)
+            tmpfile = self.tmp_path / f"tempfile_{protocol}.pkl"
 
             try:
-                with open(tmpfile, mode) as fh:
+                with open(tmpfile, "wb") as fh:
                     pickle.dump(objects, fh, protocol=protocol)
             except Exception as exc:
                 errors.append(f"pickle.dump with {protocol=} raised:\n{exc}")
                 continue
 
             try:
                 with open(tmpfile, "rb") as fh:
-                    new_objects = pickle.load(fh)
+                    unpickled_objs = pickle.load(fh)
             except Exception as exc:
                 errors.append(f"pickle.load with {protocol=} raised:\n{exc}")
                 continue
 
             # Test for equality
             if test_eq:
-                for old_obj, new_obj in zip(objects, new_objects):
-                    assert old_obj == new_obj
+                for orig, unpickled in zip(objects, unpickled_objs):
+                    assert (
+                        orig == unpickled
+                    ), f"Unpickled and original objects are unequal for {protocol=}\n{orig=}\n{unpickled=}"
 
             # Save the deserialized objects and test for equality.
-            objects_by_protocol.append(new_objects)
+            objects_by_protocol.append(unpickled_objs)
 
         if errors:
             raise ValueError("\n".join(errors))

tests/core/test_composition.py

@@ -378,9 +378,9 @@ def test_as_dict(self):
         assert d["O"] == correct_dict["O"]
 
     def test_pickle(self):
-        for c in self.comps:
-            self.serialize_with_pickle(c, test_eq=True)
-            self.serialize_with_pickle(c.to_data_dict, test_eq=True)
+        for comp in self.comps:
+            self.serialize_with_pickle(comp)
+            self.serialize_with_pickle(comp.to_data_dict)
 
     def test_to_data_dict(self):
         comp = Composition("Fe0.00009Ni0.99991")

tests/core/test_libxcfunc.py

@@ -19,7 +19,7 @@ def test_libxcfunc_api(self):
         assert xc.kind in LibxcFunc.all_kinds()
 
         # Test if object can be serialized with Pickle.
-        self.serialize_with_pickle(xc, test_eq=True)
+        self.serialize_with_pickle(xc)
 
         # Test if object supports MSONable
         self.assert_msonable(xc, test_is_subclass=False)

tests/core/test_periodic_table.py

@@ -336,13 +336,12 @@ def test_sort(self):
         assert sorted(els) == [Element.C, Element.Se]
 
     def test_pickle(self):
-        el1 = Element.Fe
-        o = pickle.dumps(el1)
-        assert el1 == pickle.loads(o)
+        pickled = pickle.dumps(Element.Fe)
+        assert Element.Fe == pickle.loads(pickled)
 
-        # Test all elements up to Uranium
-        for i in range(1, 93):
-            self.serialize_with_pickle(Element.from_Z(i), test_eq=True)
+        # Test 5 random elements
+        for idx in np.random.randint(1, 104, size=5):
+            self.serialize_with_pickle(Element.from_Z(idx))
 
     def test_print_periodic_table(self):
         Element.print_periodic_table()
@@ -385,24 +384,24 @@ def test_attr(self):
         assert self.specie4.spin == 5
 
     def test_deepcopy(self):
-        el1 = Species("Fe", 4)
-        el2 = Species("Na", 1)
-        ellist = [el1, el2]
-        assert ellist == deepcopy(ellist), "Deepcopy operation doesn't produce exact copy."
+        el1 = Species("Fe4+")
+        el2 = Species("Na1+")
+        elem_list = [el1, el2]
+        assert elem_list == deepcopy(elem_list), "Deepcopy operation doesn't produce exact copy."
 
     def test_pickle(self):
         assert self.specie1 == pickle.loads(pickle.dumps(self.specie1))
-        for i in range(1, 5):
-            self.serialize_with_pickle(getattr(self, f"specie{i}"), test_eq=True)
-        cs = Species("Cs", 1)
-        cl = Species("Cl", 1)
+        for idx in range(1, 5):
+            self.serialize_with_pickle(getattr(self, f"specie{idx}"))
+        cs = Species("Cs1+")
+        cl = Species("Cl1+")
 
-        with open("cscl.pickle", "wb") as f:
-            pickle.dump((cs, cl), f)
+        with open("cscl.pickle", "wb") as file:
+            pickle.dump((cs, cl), file)
 
-        with open("cscl.pickle", "rb") as f:
-            d = pickle.load(f)
-            assert d == (cs, cl)
+        with open("cscl.pickle", "rb") as file:
+            tup = pickle.load(file)
+            assert tup == (cs, cl)
         os.remove("cscl.pickle")
 
     def test_get_crystal_field_spin(self):

tests/core/test_units.py

@@ -268,8 +268,4 @@ def test_pickle(self):
             b = cls(10, "N bohr")
             objects = [a, b]
 
-            new_objects_from_protocol = self.serialize_with_pickle(objects)
-
-            for new_objects in new_objects_from_protocol:
-                for old_item, new_item in zip(objects, new_objects):
-                    assert str(old_item) == str(new_item)
+            self.serialize_with_pickle(objects)

tests/core/test_xcfunc.py

@@ -45,7 +45,7 @@ def test_xcfunc_api(self):
         assert ixc_11 in d
 
         # Test if object can be serialized with Pickle.
-        self.serialize_with_pickle(ixc_11, test_eq=True)
+        self.serialize_with_pickle(ixc_11)
 
         # Test if object supports MSONable
         # TODO

tests/io/abinit/test_inputs.py

@@ -122,7 +122,7 @@ def test_api(self):
         assert inp.structure == new_structure
 
         # Compatible with Pickle and MSONable?
-        self.serialize_with_pickle(inp, test_eq=False)
+        self.serialize_with_pickle(inp)
 
     def test_input_errors(self):
         """Testing typical BasicAbinitInput Error."""

tests/io/abinit/test_pseudos.py

@@ -55,7 +55,7 @@ def test_nc_pseudos(self):
                 assert pseudo.nlcc_radius >= 0.0
 
                 # Test pickle
-                self.serialize_with_pickle(pseudo, test_eq=False)
+                self.serialize_with_pickle(pseudo)
 
                 # Test MSONable
                 self.assert_msonable(pseudo)
@@ -117,16 +117,16 @@ def test_pawxml_pseudos(self):
         assert oxygen.paw_radius == approx(1.4146523028)
 
         # Test pickle
-        new_objs = self.serialize_with_pickle(oxygen, test_eq=False)
+        new_objs = self.serialize_with_pickle(oxygen)
         # Test MSONable
         self.assert_msonable(oxygen)
 
-        for o in new_objs:
-            assert o.ispaw
-            assert o.symbol == "O"
-            assert (o.Z, o.core, o.valence) == (8, 2, 6), o.Z_val == 6
+        for obj in new_objs:
+            assert obj.ispaw
+            assert obj.symbol == "O"
+            assert (obj.Z, obj.core, obj.valence) == (8, 2, 6), obj.Z_val == 6
 
-            assert o.paw_radius == approx(1.4146523028)
+            assert obj.paw_radius == approx(1.4146523028)
 
     def test_oncvpsp_pseudo_sr(self):
         """Test the ONCVPSP Ge pseudo (scalar relativistic version)."""
@@ -147,7 +147,7 @@ def test_oncvpsp_pseudo_sr(self):
         assert not ger.supports_soc
 
         # Data persistence
-        self.serialize_with_pickle(ger, test_eq=False)
+        self.serialize_with_pickle(ger)
         self.assert_msonable(ger)
 
     def test_oncvpsp_pseudo_fr(self):
@@ -157,7 +157,7 @@ def test_oncvpsp_pseudo_fr(self):
         str(pb)
 
         # Data persistence
-        self.serialize_with_pickle(pb, test_eq=False)
+        self.serialize_with_pickle(pb)
         self.assert_msonable(pb)
 
         assert pb.symbol == "Pb"