Skip to content

Commit

Permalink
Let GribOutput take a filename or file object
Browse files Browse the repository at this point in the history
  • Loading branch information
gmertes committed Jun 19, 2024
1 parent 7130984 commit b5a171d
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 b5a171d

Please sign in to comment.