From 76b299b955168bea879d5c2077c70ad91a07b0f7 Mon Sep 17 00:00:00 2001
From: Tom <thomas.hodson@ecmwf.int>
Date: Wed, 24 Jan 2024 12:20:01 +0000
Subject: [PATCH] start working on new mt keys

---
 src/ionbeam/sources/meteotracker/source.py | 17 ++++++-----------
 src/ionbeam/writers/mars_client_writer.py  | 17 +++++++----------
 2 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/src/ionbeam/sources/meteotracker/source.py b/src/ionbeam/sources/meteotracker/source.py
index 723e8aa..1b3b7a4 100644
--- a/src/ionbeam/sources/meteotracker/source.py
+++ b/src/ionbeam/sources/meteotracker/source.py
@@ -36,6 +36,11 @@
 }
 
 
+# author_patterns:
+#     Bologna:
+#         bologna_living_lab_{}: range(1,21)
+
+
 @dataclasses.dataclass
 class MeteoTrackerSource(Source):
     secrets_file: Path
@@ -46,7 +51,7 @@ class MeteoTrackerSource(Source):
     finish_after: int | None = None
     source: str = "meteotracker"
     wkt_bounding_polygon: str | None = None  # Use geojson.io to generate
-    author_filters: List[str] | None = None
+    author_patterns: dict[str, dict[str, str]] = dataclasses.field(default_factory=dict)
 
     def init(self, globals):
         super().init(globals)
@@ -75,22 +80,12 @@ def in_bounds(self, df):
         point = shapely.geometry.point.Point(lat, lon)
         return self.bounding_polygon.contains(point)
 
-    def author_filter(self, data: pd.DataFrame):
-        if self.author_filters is None:
-            return True
-        author = data["author"].iloc[0]
-        return any(re.match(filter, author) for filter in self.author_filters)
-
     def all_filters(self, path, data: pd.DataFrame):
         # bail out if the author is not one of the living labs
         if data.empty:
             logger.warning(f"{path} contains empty CSV!")
             return False
 
-        if not self.author_filter(data):
-            # logger.debug(f"Skipping {path} because author is {data['author'].iloc[0]}")
-            return False
-
         # bail out if this data is not in the target region
         if not self.in_bounds(data):
             logger.debug(f"Skipping {path} because it's not in bounds")
diff --git a/src/ionbeam/writers/mars_client_writer.py b/src/ionbeam/writers/mars_client_writer.py
index b236f2b..2962bf6 100644
--- a/src/ionbeam/writers/mars_client_writer.py
+++ b/src/ionbeam/writers/mars_client_writer.py
@@ -24,23 +24,19 @@
 import subprocess as sb
 from tqdm.notebook import tqdm
 import tempfile
+import sys
 
 
 logger = logging.getLogger(__name__)
 
 
-def write_temp_mars_request(
-    request,
-    file,
-    verb="archive",
-    quote_keys = {"source", "target", "filter"}
-):
+def write_temp_mars_request(request, file, verb="archive", quote_keys={"source", "target", "filter"}):
     keys = []
     for key, value in request.items():
         if key in quote_keys:
             keys.append(f'{key}="{value}"')
         else:
-            keys.append(f'{key}={value}')
+            keys.append(f"{key}={value}")
 
     keys = ",\n    ".join(keys)
     rendered = f"""{verb},
@@ -77,16 +73,17 @@ def process(self, message: FileMessage | FinishMessage) -> Iterable[Message]:
 
         assert message.metadata.mars_request is not None
         assert message.metadata.filepath is not None
-        
+
         request = {"database": "fdbdev", "class": "rd", "source": str(message.metadata.filepath)}
         mars_request = request | message.metadata.mars_request.as_strings()
-        
+
         logger.info(f"mars_request: {mars_request}")
 
         with tempfile.NamedTemporaryFile() as fp:
             mars_request_string = write_temp_mars_request(mars_request, file=fp.name)
             logger.debug(f"mars_request_string: {mars_request_string}")
-            run_temp_mars_request(file=fp.name)
+            # run_temp_mars_request(file=fp.name)
+            sys.exit()
 
         # TODO: the explicit mars_keys should not be necessary here.
         metadata = self.generate_metadata(message, mars_request=message.metadata.mars_request)