From 93f536fda43cf88c553b94a78e640de33ad4e1cf Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Mon, 5 Feb 2024 13:42:17 +0000
Subject: [PATCH] wip

---
 ai_models/__init__.py | 2 +-
 ai_models/__main__.py | 8 ++++++++
 ai_models/model.py    | 3 ++-
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/ai_models/__init__.py b/ai_models/__init__.py
index 22dd60b..aead1be 100644
--- a/ai_models/__init__.py
+++ b/ai_models/__init__.py
@@ -5,4 +5,4 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-__version__ = "0.3.2"
+__version__ = "0.3.3"
diff --git a/ai_models/__main__.py b/ai_models/__main__.py
index 0d96773..5e2e07f 100644
--- a/ai_models/__main__.py
+++ b/ai_models/__main__.py
@@ -36,6 +36,14 @@ def _main():
         help="Turn on debug",
     )
 
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="count",
+        default=0,
+        help="Increase verbosity",
+    )
+
     parser.add_argument(
         "--retrieve-requests",
         help=(
diff --git a/ai_models/model.py b/ai_models/model.py
index 98168c6..1ec5e1e 100644
--- a/ai_models/model.py
+++ b/ai_models/model.py
@@ -64,6 +64,7 @@ class Model:
     assets_extra_dir = None
     retrieve = {}  # Extra parameters for retrieve
     version = 1  # To be overriden in subclasses
+    grib_extra_metadata = {}  # Extra metadata for grib files
 
     param_level_ml = ([], [])  # param, level
     param_level_pl = ([], [])  # param, level
@@ -118,7 +119,7 @@ def all_fields(self):
 
     def write(self, *args, **kwargs):
         self.collect_archive_requests(
-            self.output.write(*args, **kwargs),
+            self.output.write(*args, **kwargs, **self.grib_extra_metadata),
         )
 
     def collect_archive_requests(self, written):