Skip to content

Commit

Permalink
Merge pull request #65 from ecmwf/feature/fileobj-grib-output
Browse files Browse the repository at this point in the history
Let GribOutput take a filename or file object
  • Loading branch information
b8raoult authored Jun 20, 2024
2 parents 7130984 + b5a171d commit 361109f
Showing 1 changed file with 13 additions and 2 deletions.
15 changes: 13 additions & 2 deletions src/climetlab/readers/grib/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import datetime
import logging
import re
from io import IOBase

from climetlab.decorators import normalize
from climetlab.decorators import normalize_grib_keys
Expand Down Expand Up @@ -61,9 +62,16 @@ def __getitem__(self, key):


class GribOutput:
def __init__(self, filename, split_output=False, template=None, **kwargs):
def __init__(self, file, split_output=False, template=None, **kwargs):
self._files = {}
self.filename = filename
self.fileobj = None
self.filename = None

if isinstance(file, IOBase):
self.fileobj = file
split_output = False
else:
self.filename = file

if split_output:
self.split_output = re.findall(r"\{(.*?)\}", self.filename)
Expand All @@ -80,6 +88,9 @@ def _normalize_kwargs_names(self, **kwargs):
return kwargs

def f(self, handle):
if self.fileobj:
return self.fileobj, None

if self.split_output:
path = self.filename.format(**{k: handle.get(k) for k in self.split_output})
else:
Expand Down

0 comments on commit 361109f

Please sign in to comment.