diff --git a/bin/xmet-igra-ingest b/bin/xmet-igra-ingest index 22190a5..77b5459 100755 --- a/bin/xmet-igra-ingest +++ b/bin/xmet-igra-ingest @@ -9,12 +9,11 @@ db = Database.connect(sys.argv[1]) db.execute('begin transaction') for path in sys.argv[2:]: - with open(path, 'r') as fh: - for sounding in IGRAReader.each_from_file(path): - db.add(sounding) + for sounding in IGRAReader.each_sounding_from_file(path): + db.add(sounding) - for sample in sounding.samples: - sample.sounding_id = sounding.id - db.add(sample) + for sample in sounding.samples: + sample.sounding_id = sounding.id + db.add(sample) db.commit() diff --git a/lib/xmet/igra.py b/lib/xmet/igra.py index f153920..e2ac696 100644 --- a/lib/xmet/igra.py +++ b/lib/xmet/igra.py @@ -1,6 +1,7 @@ import io import re import datetime +import zipfile import shapely from typing import Self @@ -62,6 +63,9 @@ class IGRAReader(): def __init__(self, fh: io.TextIOBase): self.fh = fh + def __del__(self): + self.fh.close() + def read_sample(self) -> SoundingSample: line = self.fh.readline() @@ -151,17 +155,30 @@ class IGRAReader(): return sounding @staticmethod - def each_from_file(path: str): - with open(path, 'r') as fh: - reader = IGRAReader(fh) + def each_sounding_from_fh(fh: io.TextIOBase): + reader = IGRAReader(fh) - while True: - sounding = reader.read() + while True: + sounding = reader.read() - if sounding is None: - break + if sounding is None: + break - yield sounding + yield sounding + + @staticmethod + def each_sounding_from_file(path: str): + if path[-4:].lower() == '.zip': + with zipfile.ZipFile(path, 'r') as z: + for member in z.infolist(): + if member.filename[-4:].lower() != '.txt': + continue + + with z.open(member.filename, 'r') as fh: + yield from IGRAReader.each_sounding_from_fh(io.TextIOWrapper(fh)) + else: + with open(path, 'r') as fh: + yield from IGRAReader.each_sounding_from_fh(fh) def cols(text: str, start: int, end: int): a = start - 1