diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f2bdbb..dca9eb0 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/batchee/harmony/service_adapter.py b/batchee/harmony/service_adapter.py index bdb9b2f..e3c47ae 100644 --- a/batchee/harmony/service_adapter.py +++ b/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() diff --git a/batchee/harmony/util.py b/batchee/harmony/util.py index f5ec011..9762232 100644 --- a/batchee/harmony/util.py +++ b/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 diff --git a/batchee/tempo_filename_parser.py b/batchee/tempo_filename_parser.py index 0ee7397..75ec0a4 100644 --- a/batchee/tempo_filename_parser.py +++ b/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 diff --git a/tests/test_filename_grouping.py b/tests/test_filename_grouping.py index d78ddb1..beecfaf 100644 --- a/tests/test_filename_grouping.py +++ b/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