Add log directory iteration helper

CHANGELOG.md

@@ -5,6 +5,7 @@ Versions follow [Semantic Versioning](https://semver.org/spec/v2.0.0.html) (`<ma
 ### Added
 * Add derived `total_accel` column to Flysight V2 IMU sensor dataframe, calculated as a vector sum of the `xyz` acceleration components
 * Add `pyflysight.log_utils.locate_log_subdir` helper for resolving child log directory from a given top-level directory
+* Add `pyflysight.log_utils.iter_log_dirs` helper for iterating through child log directories of a given top-level directory
 
 ## [v0.3.0]
 ### Changed

pyflysight/log_utils.py

@@ -1,3 +1,5 @@
+import typing as t
+from collections import abc
 from pathlib import Path
 
 import polars
@@ -61,3 +63,35 @@ def locate_log_subdir(top_dir: Path, flysight_type: FlysightType) -> Path:
         raise ValueError(f"Multiple matching log directories found. Found: {len(found_files)}")
 
     return found_files[0].parent
+
+
+class LogDir(t.NamedTuple):  # noqa: D101
+    log_dir: Path
+    flysight_type: FlysightType
+
+
+def iter_log_dirs(
+    top_dir: Path, flysight_type: FlysightType | None = None
+) -> abc.Generator[LogDir, None, None]:
+    """
+    Iterate through children of the specified top level directory & yield log directories.
+
+    A specific Flysight hardware revision can be targeted using the `flysight_type` argument; if
+    specified as `None`, both hardware types will be searched for.
+
+    NOTE: Order of yielded directories is not guaranteed.
+
+    NOTE: Directories containing trimmed log data are currently not considered.
+    """
+    possible_parents = {f.parent for f in top_dir.rglob("*.CSV")}
+
+    for p in possible_parents:
+        filenames = {f.name for f in p.glob("*") if f.is_file()}
+
+        # For now, filter out trimmed log directories
+        if "device_info.json" in filenames:
+            continue
+
+        inferred_type = classify_log_dir(p)
+        if (flysight_type is None) or (inferred_type == flysight_type):
+            yield LogDir(p, inferred_type)

tests/test_log_utils.py

@@ -4,7 +4,7 @@
 import pytest
 
 from pyflysight import FlysightType, NUMERIC_T
-from pyflysight.log_utils import classify_log_dir, get_idx, locate_log_subdir
+from pyflysight.log_utils import classify_log_dir, get_idx, iter_log_dirs, locate_log_subdir
 
 SAMPLE_DATAFRAME = polars.DataFrame(
     {
@@ -98,3 +98,32 @@ def test_locate_log_dir(tmp_path: Path, dir_structure: str, truth_parent_path: s
     truth_parent = tmp_path / truth_parent_path
     assert locate_log_subdir(tmp_path, FlysightType.VERSION_1) == truth_parent
     assert locate_log_subdir(tmp_path, FlysightType.VERSION_2) == truth_parent
+
+
+ITER_DIR_STRUCTURE = {
+    "24-04-20": ("04-20-00.CSV",),
+    "12-34-00": ("RAW.UBX", "SENSOR.CSV", "TRACK.CSV"),
+    "abc123": ("BARO.CSV", "IMU.CSV", "TRACK.CSV", "device_info.json"),
+}
+
+ITER_DIR_TEST_CASES = (
+    (None, {"24-04-20", "12-34-00"}),
+    (FlysightType.VERSION_1, {"24-04-20"}),
+    (FlysightType.VERSION_2, {"12-34-00"}),
+)
+
+
+@pytest.mark.parametrize(("hw_type", "truth_parent_dirnames"), ITER_DIR_TEST_CASES)
+def test_iter_dir(
+    tmp_path: Path, hw_type: FlysightType | None, truth_parent_dirnames: set[str]
+) -> None:
+    for log_dirname, filenames in ITER_DIR_STRUCTURE.items():
+        log_dir = tmp_path / log_dirname
+        log_dir.mkdir()
+
+        for name in filenames:
+            (log_dir / name).touch()
+
+    found_dirs = iter_log_dirs(top_dir=tmp_path, flysight_type=hw_type)
+    found_dirnames = {ld.log_dir.name for ld in found_dirs}
+    assert found_dirnames == truth_parent_dirnames