diff --git a/README.md b/README.md index ab57a35..959f88a 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,17 @@ dd = SyncIfcbDataDirectory('/path/to/ifcb/data') for fileset in dd.list(): print(fileset['pid']) # e.g. 'D20221227T093138_IFCB127' +# Filter the listing by timestamp range and/or instrument. +# The range is half-open [start_time, end_time); naive datetimes are UTC. +from datetime import datetime, timezone + +for fileset in dd.list( + start_time=datetime(2022, 12, 1, tzinfo=timezone.utc), + end_time=datetime(2023, 1, 1, tzinfo=timezone.utc), + instrument=127, # int, or an iterable like [127, 130] +): + print(fileset['pid']) + # Check if a specific bin exists dd.exists('D20221227T093138_IFCB127') diff --git a/src/ifcbkit/__init__.py b/src/ifcbkit/__init__.py index a679cf4..bfde9cb 100644 --- a/src/ifcbkit/__init__.py +++ b/src/ifcbkit/__init__.py @@ -43,6 +43,7 @@ # Fileset discovery from .fileset import ( validate_path, + make_fileset_filter, async_list_filesets, sync_list_filesets, async_list_data_dirs, diff --git a/src/ifcbkit/fileset.py b/src/ifcbkit/fileset.py index 3b911c6..bfd552f 100644 --- a/src/ifcbkit/fileset.py +++ b/src/ifcbkit/fileset.py @@ -10,11 +10,12 @@ import asyncio import os +from datetime import timezone import aiofiles.os as aios import aiofiles.ospath as aiopath -from .identifiers import add_target, parse_roi_id +from .identifiers import add_target, parse_roi_id, bin_timestamp, bin_instrument_id DEFAULT_EXCLUDE = ['skip', 'beads'] @@ -92,6 +93,65 @@ def validate_path( return True +# --- Fileset filtering: timestamp range / instrument --- + +def _normalize_filter_time(dt): + """Assume UTC for naive datetimes; leave aware datetimes untouched.""" + if dt is None: + return None + if dt.tzinfo is None: + return dt.replace(tzinfo=timezone.utc) + return dt + + +def make_fileset_filter(start_time=None, end_time=None, instrument=None): + """ + Build a predicate ``(basename) -> bool`` for filtering filesets. + + :param start_time: inclusive lower bound (datetime); naive treated as UTC + :param end_time: exclusive upper bound (datetime); naive treated as UTC + :param instrument: instrument ID (int or str) or iterable of instrument IDs + :returns: a predicate accepting a bin ID basename, or None if no filter is + active. Basenames that fail to parse are excluded when any filter is set. + :raises ValueError: if ``instrument`` cannot be coerced to int(s) + + The timestamp range is half-open ``[start_time, end_time)``. + """ + if start_time is None and end_time is None and instrument is None: + return None + + start = _normalize_filter_time(start_time) + end = _normalize_filter_time(end_time) + + if instrument is None: + instruments = None + else: + # Coerce scalars (int, or str like "127") and iterables of IDs to a set + # of ints. Anything that won't coerce is a ValueError. + scalars = (int, str) if not isinstance(instrument, bool) else () + values = [instrument] if isinstance(instrument, scalars) else instrument + try: + instruments = {int(v) for v in values} + except (TypeError, ValueError) as e: + raise ValueError(f'invalid instrument filter: {instrument!r}') from e + + def _pred(basename): + try: + if instruments is not None and bin_instrument_id(basename) not in instruments: + return False + if start is not None or end is not None: + ts = bin_timestamp(basename) + if start is not None and ts < start: + return False + if end is not None and ts >= end: + return False + except ValueError: + return False + return True + + return _pred + + # --- Internal helpers: directory entry splitting --- async def _async_split_dir_entries(dirpath, *, exclude=DEFAULT_EXCLUDE, sort=True, reverse=False): @@ -146,6 +206,9 @@ async def async_list_filesets( validate=True, require_adc=True, require_roi=True, + start_time=None, + end_time=None, + instrument=None, ): """ Async generator yielding (dp, basename) for each .hdr/.adc/(.roi) fileset found. @@ -157,10 +220,15 @@ async def async_list_filesets( :param validate: whether to validate paths :param require_adc: require .adc file presence :param require_roi: require .roi file presence + :param start_time: inclusive lower bound on bin timestamp (datetime, UTC if naive) + :param end_time: exclusive upper bound on bin timestamp (datetime, UTC if naive) + :param instrument: instrument ID (int) or iterable of instrument IDs to keep """ if not set(exclude).isdisjoint(set(include)): raise ValueError('include and exclude must be disjoint') + fs_filter = make_fileset_filter(start_time, end_time, instrument) + stack = [dirpath] while stack: dp = stack.pop() @@ -182,6 +250,8 @@ async def async_list_filesets( reldir = dp[len(dirpath) + 1:] if not validate_path(os.path.join(reldir, basename), include=include, exclude=exclude): continue + if fs_filter is not None and not fs_filter(basename): + continue yield dp, basename @@ -193,6 +263,9 @@ def sync_list_filesets( validate=True, require_adc=True, require_roi=True, + start_time=None, + end_time=None, + instrument=None, ): """ Sync generator yielding (dp, basename) for each .hdr/.adc/(.roi) fileset found. @@ -204,10 +277,15 @@ def sync_list_filesets( :param validate: whether to validate paths :param require_adc: require .adc file presence :param require_roi: require .roi file presence + :param start_time: inclusive lower bound on bin timestamp (datetime, UTC if naive) + :param end_time: exclusive upper bound on bin timestamp (datetime, UTC if naive) + :param instrument: instrument ID (int) or iterable of instrument IDs to keep """ if not set(exclude).isdisjoint(set(include)): raise ValueError('include and exclude must be disjoint') + fs_filter = make_fileset_filter(start_time, end_time, instrument) + stack = [dirpath] while stack: dp = stack.pop() @@ -229,6 +307,8 @@ def sync_list_filesets( reldir = dp[len(dirpath) + 1:] if not validate_path(os.path.join(reldir, basename), include=include, exclude=exclude): continue + if fs_filter is not None and not fs_filter(basename): + continue yield dp, basename @@ -441,12 +521,18 @@ def paths(self, pid): 'roi': fs + '.roi' if self.require_roi else None, } - def list(self): - """Yield dicts of {pid, hdr, adc, roi} for all filesets.""" + def list(self, start_time=None, end_time=None, instrument=None): + """Yield dicts of {pid, hdr, adc, roi} for all filesets. + + :param start_time: inclusive lower bound on bin timestamp (UTC if naive) + :param end_time: exclusive upper bound on bin timestamp (UTC if naive) + :param instrument: instrument ID (int) or iterable of instrument IDs + """ for dp, bn in sync_list_filesets( self.root_path, exclude=self.exclude, include=self.include, require_adc=self.require_adc, require_roi=self.require_roi, + start_time=start_time, end_time=end_time, instrument=instrument, ): yield { 'pid': bn, @@ -571,12 +657,18 @@ async def paths(self, pid): 'roi': fs + '.roi' if self.require_roi else None, } - async def list(self): - """Async generator yielding dicts of {pid, hdr, adc, roi} for all filesets.""" + async def list(self, start_time=None, end_time=None, instrument=None): + """Async generator yielding dicts of {pid, hdr, adc, roi} for all filesets. + + :param start_time: inclusive lower bound on bin timestamp (UTC if naive) + :param end_time: exclusive upper bound on bin timestamp (UTC if naive) + :param instrument: instrument ID (int) or iterable of instrument IDs + """ async for dp, bn in async_list_filesets( self.root_path, exclude=self.exclude, include=self.include, require_adc=self.require_adc, require_roi=self.require_roi, + start_time=start_time, end_time=end_time, instrument=instrument, ): adc = None if self.require_adc: diff --git a/tests/test_fileset.py b/tests/test_fileset.py index 851ad22..4eade36 100644 --- a/tests/test_fileset.py +++ b/tests/test_fileset.py @@ -3,10 +3,13 @@ import asyncio import os import shutil +from datetime import datetime, timezone import pytest -from ifcbkit.fileset import SyncIfcbDataDirectory, AsyncIfcbDataDirectory +from ifcbkit.fileset import ( + SyncIfcbDataDirectory, AsyncIfcbDataDirectory, make_fileset_filter, +) PID = 'D20170426T164105_IFCB009' DAY = 'D20170426' @@ -171,3 +174,94 @@ def test_istyle_corrected_adc_content_is_used(tmp_path): corrected = SyncIfcbDataDirectory(str(root)).read_images(I_PID) assert dropped in baseline assert dropped not in corrected + + +# --- list() filtering: timestamp range / instrument --- + +# Bins spanning two days, two instruments. Timestamps parsed from the PIDs. +BIN_A = 'D20200101T000000_IFCB100' # 2020-01-01, instr 100 +BIN_B = 'D20200102T120000_IFCB100' # 2020-01-02, instr 100 +BIN_C = 'D20200103T000000_IFCB200' # 2020-01-03, instr 200 + + +def _make_multi_bin_root(tmp_path): + root = tmp_path / 'data' + for pid in (BIN_A, BIN_B, BIN_C): + _make_fileset(str(root / pid[:9]), pid) + return root + + +def _pids(entries): + return sorted(e['pid'] for e in entries) + + +def test_make_fileset_filter_none_when_no_args(): + assert make_fileset_filter() is None + + +def test_list_no_filter_returns_all(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + assert _pids(dd.list()) == [BIN_A, BIN_B, BIN_C] + + +def test_list_filter_by_instrument_int(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + assert _pids(dd.list(instrument=200)) == [BIN_C] + + +def test_list_filter_by_instrument_iterable(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + assert _pids(dd.list(instrument=[100, 200])) == [BIN_A, BIN_B, BIN_C] + + +def test_list_filter_by_instrument_str(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + assert _pids(dd.list(instrument='200')) == [BIN_C] + + +def test_list_filter_by_instrument_str_iterable(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + assert _pids(dd.list(instrument=['100', '200'])) == [BIN_A, BIN_B, BIN_C] + + +def test_make_fileset_filter_bad_instrument_raises(): + with pytest.raises(ValueError): + make_fileset_filter(instrument='not-an-int') + with pytest.raises(ValueError): + make_fileset_filter(instrument=['100', 'bad']) + with pytest.raises(ValueError): + make_fileset_filter(instrument=True) + + +def test_list_filter_start_time_inclusive(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + start = datetime(2020, 1, 2, 12, 0, 0, tzinfo=timezone.utc) + assert _pids(dd.list(start_time=start)) == [BIN_B, BIN_C] + + +def test_list_filter_end_time_exclusive(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + end = datetime(2020, 1, 3, 0, 0, 0, tzinfo=timezone.utc) + assert _pids(dd.list(end_time=end)) == [BIN_A, BIN_B] + + +def test_list_filter_range_and_instrument(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + start = datetime(2020, 1, 1, tzinfo=timezone.utc) + end = datetime(2020, 1, 3, tzinfo=timezone.utc) + assert _pids(dd.list(start_time=start, end_time=end, instrument=100)) == [BIN_A, BIN_B] + + +def test_list_filter_naive_datetime_treated_utc(tmp_path): + dd = SyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + start = datetime(2020, 1, 3) # naive -> UTC + assert _pids(dd.list(start_time=start)) == [BIN_C] + + +def test_async_list_filter(tmp_path): + dd = AsyncIfcbDataDirectory(str(_make_multi_bin_root(tmp_path))) + + async def _collect(): + return [e async for e in dd.list(instrument=100)] + + assert _pids(asyncio.run(_collect())) == [BIN_A, BIN_B]