Do not set items in dictionary that would be None

vasp-dev · Aug 30, 2023 · 95f92d9 · 95f92d9
1 parent 713b66c
commit 95f92d9
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 23 deletions.
diff --git a/src/py4vasp/_data/contcar.py b/src/py4vasp/_data/contcar.py
@@ -14,9 +14,9 @@ def to_dict(self):
         return {
             **self._structure().read(),
             "system": convert.text_to_string(self._raw_data.system),
-            "selective_dynamics": self._read_selective_dynamics(),
-            "lattice_velocities": self._read_lattice_velocities(),
-            "ion_velocities": self._read_ion_velocities(),
+            **self._read("selective_dynamics"),
+            **self._read("lattice_velocities"),
+            **self._read("ion_velocities"),
         }
 
     def _structure(self):
@@ -31,20 +31,6 @@ def _structure(self):
         )
         return data.Structure.from_data(raw_structure)
 
-    def _read_selective_dynamics(self):
-        if self._raw_data.selective_dynamics.is_none():
-            return None
-        else:
-            return self._raw_data.selective_dynamics[:]
-
-    def _read_lattice_velocities(self):
-        if self._raw_data.lattice_velocities.is_none():
-            return None
-        else:
-            return self._raw_data.lattice_velocities[:]
-
-    def _read_ion_velocities(self):
-        if self._raw_data.ion_velocities.is_none():
-            return None
-        else:
-            return self._raw_data.ion_velocities[:]
+    def _read(self, key):
+        data = getattr(self._raw_data, key)
+        return {key: data} if not data.is_none() else {}
diff --git a/tests/data/test_contcar.py b/tests/data/test_contcar.py
@@ -21,6 +21,18 @@ def CONTCAR(raw_data, request):
     return contcar
 
 
+class OptionalOutputCheck:
+    def __init__(self, dict_, Assert):
+        self.dict_ = dict_
+        self.Assert = Assert
+
+    def element_agrees(self, key, reference):
+        if reference.is_none():
+            assert key not in self.dict_
+        else:
+            self.Assert.allclose(self.dict_[key], reference)
+
+
 def test_read(CONTCAR, Assert):
     actual = CONTCAR.read()
     expected = CONTCAR.ref.structure.read()
@@ -29,6 +41,7 @@ def test_read(CONTCAR, Assert):
     assert actual["elements"] == expected["elements"]
     assert actual["names"] == expected["names"]
     assert actual["system"] == CONTCAR.ref.system
-    Assert.allclose(actual["selective_dynamics"], CONTCAR.ref.selective_dynamics)
-    Assert.allclose(actual["lattice_velocities"], CONTCAR.ref.lattice_velocities)
-    Assert.allclose(actual["ion_velocities"], CONTCAR.ref.ion_velocities)
+    check = OptionalOutputCheck(actual, Assert)
+    check.element_agrees("selective_dynamics", CONTCAR.ref.selective_dynamics)
+    check.element_agrees("lattice_velocities", CONTCAR.ref.lattice_velocities)
+    check.element_agrees("ion_velocities", CONTCAR.ref.ion_velocities)