Skip to content

Commit

Permalink
start working on new mt keys
Browse files Browse the repository at this point in the history
  • Loading branch information
TomHodson committed Jan 24, 2024
1 parent a0d9c91 commit 76b299b
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 21 deletions.
17 changes: 6 additions & 11 deletions src/ionbeam/sources/meteotracker/source.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,11 @@
}


# author_patterns:
# Bologna:
# bologna_living_lab_{}: range(1,21)


@dataclasses.dataclass
class MeteoTrackerSource(Source):
secrets_file: Path
Expand All @@ -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)
Expand Down Expand Up @@ -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")
Expand Down
17 changes: 7 additions & 10 deletions src/ionbeam/writers/mars_client_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 76b299b

Please sign in to comment.