TRADE-1167: Refactored grouping logic into separate function for reuse

CHANGELOG.md

@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Refactored grouping logic into separate function for reuse ([#215](https://github.com/nasa/batchee/issues/215)) ([**kecunning**](https://github.com/kecunning))
+
 - Implemented new streamlined release workflow. ([#213](https://github.com/nasa/batchee/issues/213))([**@ank1m**](https://github.com/ank1m))
 - Switched dependency management and packaging from Poetry to uv for faster installs and simplified workflow. ([#182](https://github.com/nasa/batchee/issues/182))([**@ank1m**](https://github.com/ank1m))
 

batchee/harmony/service_adapter.py

@@ -36,6 +36,7 @@
     _get_item_url,
     _get_netcdf_urls,
     _get_output_date_range,
+    _group_batch_indices,
 )
 from batchee.tempo_filename_parser import get_batch_indices
 
@@ -100,14 +101,12 @@ def process_catalog(self, catalog: pystac.Catalog) -> list[pystac.Catalog]:
             self.logger.info(f"batch_indices==={batch_indices}.")
 
             # --- Construct a dictionary with a separate key for each batch ---
-            grouped: dict[int, list[Item]] = {}
-            for k, v in zip(batch_indices, items, strict=False):
-                grouped.setdefault(k, []).append(v)
+            grouped_batches = _group_batch_indices(batch_indices, items)
 
             # --- Construct a list of STAC Catalogs (which represent each TEMPO scan),
             #   and each Catalog holds multiple Items (which represent each granule).
             catalogs = []
-            for batch_id, batch_items in grouped.items():
+            for batch_id, batch_items in grouped_batches.items():
                 self.logger.info(f"constructing new pystac.Catalog for batch_id==={batch_id}.")
                 # Initialize a new, empty Catalog
                 batch_catalog = catalog.clone()

batchee/harmony/util.py

@@ -27,6 +27,7 @@
 """Misc utility functions"""
 
 from datetime import datetime
+from typing import Any
 
 from pystac import Asset, Item
 
@@ -129,3 +130,14 @@ def _get_item_date_range(item: Item) -> tuple[datetime, datetime]:
         end_datetime = item.datetime
 
     return start_datetime, end_datetime
+
+
+def _group_batch_indices(batch_indices: list, items: list) -> dict[int, list]:
+    """A helper funtion that groups batches and constructs a dictionary
+    with a separate key for each batch
+    """
+    grouped: dict[int, Any] = {}
+    for k, v in zip(batch_indices, items, strict=False):
+        grouped.setdefault(k, []).append(v)
+
+    return grouped

batchee/tempo_filename_parser.py

@@ -32,6 +32,8 @@
 from datetime import datetime
 from zoneinfo import ZoneInfo
 
+from batchee.harmony.util import _group_batch_indices
+
 default_logger = logging.getLogger(__name__)
 
 tempo_granule_filename_pattern = re.compile(
@@ -144,11 +146,9 @@ def main() -> list[list[str]]:
     unique_category_indices: list[int] = sorted(set(batch_indices), key=batch_indices.index)
     logging.info(f"batch_indices = {batch_indices}")
 
-    # --- Construct a STAC object based on the batch indices ---
-    grouped: dict[int, list[str]] = {}
-    for k, v in zip(batch_indices, input_filenames, strict=False):
-        grouped.setdefault(k, []).append(v)
-    grouped_names: list[list[str]] = [grouped[k] for k in unique_category_indices]
+    # --- Construct a dictionary with a separate key for each batch ---
+    grouped_batches = _group_batch_indices(batch_indices, input_filenames)
+    grouped_names: list[list[str]] = [grouped_batches[k] for k in unique_category_indices]
 
     return grouped_names
 

tests/test_filename_grouping.py

@@ -2,7 +2,11 @@
 from unittest.mock import patch
 
 import batchee.tempo_filename_parser
-from batchee.tempo_filename_parser import get_batch_indices, get_day_in_us_central
+from batchee.tempo_filename_parser import (
+    _group_batch_indices,
+    get_batch_indices,
+    get_day_in_us_central,
+)
 
 example_filenames = [
     "TEMPO_NO2_L2_V03_20240731T235252Z_S016G04.nc",
@@ -28,6 +32,14 @@
     "TEMPO_NO2_L2_NRT_V02_20250712T170552Z_S008G09.nc",
 ]
 
+example_grouping = {
+    0: [
+        "TEMPO_NO2_L2_V04_20230827T002226Z_S011G07.nc",
+        "TEMPO_NO2_L2_V04_20230827T002839Z_S011G08.nc",
+    ],
+    1: ["TEMPO_NO2_L2_V04_20230831T200444Z_S012G07.nc"],
+}
+
 
 def test_timezone_conversion():
     utcdates = [filename.split("_")[4] for filename in example_filenames]
@@ -76,3 +88,14 @@ def test_invalid_filenames():
     invalid_filenames = ["invalid.nc", "TEMPO_NO2_L2_V03_20240731.nc", "random_file.txt"]
     results = get_batch_indices(invalid_filenames)
     assert results == []
+
+
+def test_file_grouping():
+    filenames = [
+        "TEMPO_NO2_L2_V04_20230827T002226Z_S011G07.nc",
+        "TEMPO_NO2_L2_V04_20230827T002839Z_S011G08.nc",
+        "TEMPO_NO2_L2_V04_20230831T200444Z_S012G07.nc",
+    ]
+    batch_indices = get_batch_indices(filenames)
+    grouped_batches = _group_batch_indices(batch_indices, filenames)
+    assert grouped_batches == example_grouping