better error handeling

Author: Han Wang

dpdata/lmdb/format.py

@@ -13,6 +13,18 @@
 m.patch()
 
 
+class LMDBError(Exception):
+    """Base class for LMDB errors."""
+
+
+class LMDBMetadataError(LMDBError):
+    """Metadata not found in LMDB."""
+
+
+class LMDBFrameError(LMDBError):
+    """Frame data not found in LMDB."""
+
+
 class LMDBFormat(Format):
     """
     Class for handling the LMDB format, which stores atomic configurations in a
@@ -194,7 +206,7 @@ def from_multi_systems(self, file_name, map_size=1000000000, **kwargs):
             with env.begin() as txn:
                 metadata_packed = txn.get(b"__metadata__")
                 if metadata_packed is None:
-                    raise KeyError("LMDB database does not contain metadata.")
+                    raise LMDBMetadataError("LMDB database does not contain metadata.")
                 metadata = msgpack.unpackb(metadata_packed, raw=False)
                 frame_idx_fmt = metadata.get("frame_idx_fmt", "012d")
 
@@ -209,7 +221,7 @@ def from_multi_systems(self, file_name, map_size=1000000000, **kwargs):
                         key = f"{i:{frame_idx_fmt}}".encode("ascii")
                         value = txn.get(key)
                         if value is None:
-                            raise KeyError(f"Frame data not found for key: {key}")
+                            raise LMDBFrameError(f"Frame data not found for key: {key}")
                         frame_data = msgpack.unpackb(value, raw=False)
                         system_frames.append(frame_data)
 

tests/test_lmdb.py

@@ -17,6 +17,7 @@
 )
 from context import dpdata
 
+from dpdata.lmdb.format import LMDBFrameError, LMDBMetadataError
 from dpdata.plugins.lmdb import LMDBFormat
 
 
@@ -128,13 +129,13 @@ def test_load_missing_metadata(self):
         ).close()  # Creates empty LMDB environment
 
         with self.assertRaisesRegex(
-            KeyError, "LMDB database does not contain metadata."
+            LMDBMetadataError, "LMDB database does not contain metadata."
         ):
             LMDBFormat().from_multi_systems(self.lmdb_path_missing_metadata)
 
     def test_load_missing_frame_data(self):
         with self.assertRaisesRegex(
-            KeyError, "Frame data not found for key: b'000000000000'"
+            LMDBFrameError, "Frame data not found for key: b'000000000000'"
         ):
             LMDBFormat().from_multi_systems(self.lmdb_path_missing_frame)