From fb1fd116b2b99b4def48bd8c1d4a876ac53761d4 Mon Sep 17 00:00:00 2001
From: Gert Mertes <hey@gertmertes.be>
Date: Tue, 18 Jun 2024 15:49:03 +0000
Subject: [PATCH] Refactor FileOutput to inherit from FileOutputBase

---
 src/ai_models/outputs/__init__.py | 36 +++++++++++++++++++------------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/src/ai_models/outputs/__init__.py b/src/ai_models/outputs/__init__.py
index cba9f8a..20874b2 100644
--- a/src/ai_models/outputs/__init__.py
+++ b/src/ai_models/outputs/__init__.py
@@ -25,29 +25,31 @@ def finalise(self, *args, **kwargs):
         pass
 
 
-class FileOutput(Output):
+class FileOutputBase(Output):
     def __init__(self, owner, path, metadata, **kwargs):
         self._first = True
         metadata.setdefault("stream", "oper")
         metadata.setdefault("expver", owner.expver)
         metadata.setdefault("class", "ml")
 
-        LOG.info("Writing results to %s.", path)
         self.path = path
         self.owner = owner
         self.metadata = metadata
 
     @cached_property
-    def output(self):
-
+    def grib_keys(self):
         edition = self.metadata.pop("edition", 2)
 
-        self.grib_keys = dict(
+        _grib_keys = dict(
             edition=edition,
             generatingProcessIdentifier=self.owner.version,
         )
-        self.grib_keys.update(self.metadata)
+        _grib_keys.update(self.metadata)
+
+        return _grib_keys
 
+    @cached_property
+    def output(self):
         return cml.new_grib_output(
             self.path,
             split_output=True,
@@ -91,6 +93,20 @@ def write(self, data, *args, check=False, **kwargs):
         return handle, path
 
 
+class FileOutput(FileOutputBase):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        LOG.info("Writing results to %s", self.path)
+
+
+class NoneOutput(Output):
+    def __init__(self, *args, **kwargs):
+        LOG.info("Results will not be written.")
+
+    def write(self, *args, **kwargs):
+        pass
+
+
 class HindcastReLabel:
     def __init__(self, owner, output, hindcast_reference_year=None, hindcast_reference_date=None, **kwargs):
         self.owner = owner
@@ -151,14 +167,6 @@ def write(self, *args, **kwargs):
         return self.output.write(*args, **kwargs)
 
 
-class NoneOutput(Output):
-    def __init__(self, *args, **kwargs):
-        LOG.info("Results will not be written.")
-
-    def write(self, *args, **kwargs):
-        pass
-
-
 def get_output(name, owner, *args, **kwargs):
     result = available_outputs()[name].load()(owner, *args, **kwargs)
     if kwargs.get("hindcast_reference_year") is not None or kwargs.get("hindcast_reference_date") is not None: