diff --git a/lib/nexrad/archive.py b/lib/nexrad/archive.py index ebd1e32..ea76d34 100644 --- a/lib/nexrad/archive.py +++ b/lib/nexrad/archive.py @@ -3,6 +3,14 @@ import re from nexrad.s3 import S3Bucket +class ArchiveDateError(Exception): + def __init__(self, supplied: str, missing: str): + self.supplied = supplied + self.missing = missing + + def __str__(self): + return "Date {self.supplied} was supplied, but required {self.missing} is missing" + class Archive(): path: str bucket: S3Bucket @@ -28,26 +36,44 @@ class Archive(): RE_CALL = re.compile(r'^[A-Z]{4}$') RE_PRODUCT = re.compile(r'^([A-Z]{4})(\d{4})(\d{2})(\d{2})_(\d{2})(\d{2})(\d{2})') - def each_downloaded_key(self): + def each_downloaded_key(self, + year: int=None, + month: int=None, + day: int=None): parts = [self.path] - for year in os.scandir(os.path.join(*parts)): - if not (year.is_dir() and self.RE_YEAR.match(year.name)): + if day is not None and month is None: + raise ArchiveDateError('day', 'month') + + if month is not None and year is None: + raise ArchiveDateError('month', 'year') + + for cur_year in os.scandir(os.path.join(*parts)): + if not (cur_year.is_dir() and self.RE_YEAR.match(cur_year.name)): continue - parts.append(year.name) + if year is not None and int(cur_year.name) != year: + continue - for month in os.scandir(os.path.join(*parts)): - if not (month.is_dir() and self.RE_MONTH_DAY.match(month.name)): + parts.append(cur_year.name) + + for cur_month in os.scandir(os.path.join(*parts)): + if not (cur_month.is_dir() and self.RE_MONTH_DAY.match(cur_month.name)): continue - parts.append(month.name) + if month is not None and int(cur_month.name) != month: + continue - for day in os.scandir(os.path.join(*parts)): - if not (day.is_dir() and self.RE_MONTH_DAY.match(day.name)): + parts.append(cur_month.name) + + for cur_day in os.scandir(os.path.join(*parts)): + if not (cur_day.is_dir() and self.RE_MONTH_DAY.match(cur_day.name)): continue - parts.append(day.name) + if day is not None and int(cur_day.name) != day: + continue + + parts.append(cur_day.name) for call in os.scandir(os.path.join(*parts)): if not (call.is_dir() and self.RE_CALL.match(call.name)):