clean up a little

Author: czaloom

src/valor_lite/cache.py

@@ -8,6 +8,7 @@
 import numpy as np
 import pyarrow as pa
 import pyarrow.dataset as ds
+import pyarrow.lib as pl
 import pyarrow.parquet as pq
 
 
@@ -40,6 +41,27 @@ def to_arrow(self):
                 return pa.timestamp("us")
 
 
+def convert_type_mapping_to_schema(
+    type_mapping: dict[str, DataType] | None
+) -> list[tuple[str, pl.DataType]]:
+    """
+    Convert type mapping to a pyarrow schema input.
+
+    Parameters
+    ----------
+    type_mapping : dict[str, DataType] | None
+        A map from string key to datatype. Treats input of `None` as empty mapping.
+
+    Returns
+    -------
+    list[tuple[str, pyarrow.lib.DataType]]
+        A list of field name, field type pairs that can be used as input to pyarrow.schema.
+    """
+    if not type_mapping:
+        return []
+    return [(k, v.to_arrow()) for k, v in type_mapping.items()]
+
+
 class CacheReader:
     def __init__(self, where: str | Path):
         self._dir = Path(where)

src/valor_lite/object_detection/loader.py

@@ -6,7 +6,11 @@
 from numpy.typing import NDArray
 from tqdm import tqdm
 
-from valor_lite.cache import CacheWriter, DataType
+from valor_lite.cache import (
+    CacheWriter,
+    DataType,
+    convert_type_mapping_to_schema,
+)
 from valor_lite.exceptions import EmptyCacheError
 from valor_lite.object_detection.annotation import (
     Bitmask,
@@ -56,20 +60,14 @@ def __init__(
             }
             json.dump(types, f, indent=2)
 
-        datum_metadata_schema = (
-            [(k, v.to_arrow()) for k, v in datum_metadata_types.items()]
-            if datum_metadata_types
-            else []
+        datum_metadata_schema = convert_type_mapping_to_schema(
+            datum_metadata_types
         )
-        groundtruth_metadata_schema = (
-            [(k, v.to_arrow()) for k, v in groundtruth_metadata_types.items()]
-            if groundtruth_metadata_types
-            else []
+        groundtruth_metadata_schema = convert_type_mapping_to_schema(
+            groundtruth_metadata_types
         )
-        prediction_metadata_schema = (
-            [(k, v.to_arrow()) for k, v in prediction_metadata_types.items()]
-            if prediction_metadata_types
-            else []
+        prediction_metadata_schema = convert_type_mapping_to_schema(
+            prediction_metadata_types
         )
 
         self._null_gt_metadata = {