Split encoding from writing in grib output (#409)

* Split encoding from writing in grib output

Author: sandorkertesz

src/earthkit/data/readers/grib/output.py

@@ -10,15 +10,14 @@
 import datetime
 import logging
 import re
+from io import IOBase
 
 from earthkit.data.decorators import normalize
 from earthkit.data.decorators import normalize_grib_keys
 from earthkit.data.utils.humanize import list_to_human
 
 LOG = logging.getLogger(__name__)
 
-# Make sure the
-
 ACCUMULATIONS = {("tp", 2): {"productDefinitionTemplateNumber": 8}}
 
 _ORDER = (
@@ -59,15 +58,8 @@ def __getitem__(self, key):
         return self.handle.get(key, default=None)
 
 
-class GribOutput:
-    def __init__(self, filename, split_output=False, template=None, **kwargs):
-        self._files = {}
-        self.filename = filename
-
-        if split_output:
-            self.split_output = re.findall(r"\{(.*?)\}", self.filename)
-        else:
-            self.split_output = None
+class GribCoder:
+    def __init__(self, template=None, **kwargs):
 
         self.template = template
         self._bbox = {}
@@ -78,22 +70,13 @@ def __init__(self, filename, split_output=False, template=None, **kwargs):
     def _normalize_kwargs_names(self, **kwargs):
         return kwargs
 
-    def f(self, handle):
-        if self.split_output:
-            path = self.filename.format(**{k: handle.get(k) for k in self.split_output})
-        else:
-            path = self.filename
-
-        if path not in self._files:
-            self._files[path] = open(path, "wb")
-        return self._files[path], path
-
-    def write(
+    def encode(
         self,
         values,
         check_nans=False,
         metadata={},
         template=None,
+        return_bytes=False,
         **kwargs,
     ):
         # Make a copy as we may modify it
@@ -148,20 +131,10 @@ def write(
         if values is not None:
             handle.set_values(values)
 
-        file, path = self.f(handle)
-        handle.write(file)
+        if return_bytes:
+            return handle.get_message()
 
-        return handle, path
-
-    def close(self):
-        for f in self._files.values():
-            f.close()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, trace):
-        self.close()
+        return handle
 
     def update_metadata(self, handle, metadata, compulsory):
         # TODO: revisit that logic
@@ -376,5 +349,73 @@ def _gg_field(self, values, metadata):
             return f"reduced_gg_{levtype}_{N}_grib{edition}"
 
 
+class GribOutput:
+    def __init__(self, file, split_output=False, template=None, **kwargs):
+        self._files = {}
+        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)
+        else:
+            self.split_output = None
+
+        self._coder = GribCoder(template=template, **kwargs)
+
+    def close(self):
+        for f in self._files.values():
+            f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, trace):
+        self.close()
+
+    def write(
+        self,
+        values,
+        check_nans=False,
+        metadata={},
+        template=None,
+        **kwargs,
+    ):
+        handle = self._coder.encode(
+            values,
+            check_nans=check_nans,
+            metadata=metadata,
+            template=template,
+            **kwargs,
+        )
+
+        file, path = self.f(handle)
+        handle.write(file)
+
+        return handle, path
+
+    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:
+            path = self.filename
+
+        if path not in self._files:
+            self._files[path] = open(path, "wb")
+        return self._files[path], path
+
+
 def new_grib_output(*args, **kwargs):
     return GribOutput(*args, **kwargs)
+
+
+def new_grib_coder(*args, **kwargs):
+    return GribCoder(*args, **kwargs)