from dataclasses import dataclass from datetime import datetime from enum import Enum def _json(value): try: return value.json() except AttributeError: return value def _convert_notation(unit): return unit.replace("-1", "^-1") class _StrEnum(Enum): def __str__(self): return self.value def json(self): return self.value class Model(_StrEnum): AROME = "arome" GEOS5 = "geos5" GFS = "gfs" GFSWAVE = "gfsWave" ICONEU = "iconEu" NAMALASKA = "namAlaska" NAMCONUS = "namConus" NAMHAWAII = "namHawaii" class Level(_StrEnum): SURFACE = "surface" H1000 = "1000h" H950 = "950h" H925 = "925h" H900 = "900h" H850 = "850h" H800 = "800h" H700 = "700h" H600 = "600h" H500 = "500h" H400 = "400h" H300 = "300h" H200 = "200h" H150 = "150h" @dataclass class Request: key: str lat: float lon: float model: Model parameters: list = None levels: list = None def json(self): body = { 'key': self.key, 'lat': self.lat, 'lon': self.lon, 'model': _json(self.model), 'parameters': self.parameters or [], } if self.levels: body['levels'] = [_json(x) for x in self.levels] return body class EntryView: def __init__(self, response, index=0): self._response = response self._index = index def __iter__(self): return self def __next__(self): self._index += 1 if not self: raise StopIteration return self def __bool__(self): return self._index < len(self) def __len__(self): return len(self._response) @property def timestamp(self): return self._response.timestamps[self._index] def __getitem__(self, key): return self._response.samples[key][self._index] class Response: _INTERNAL_FIELDS = ('ts', 'units', 'warning') def __init__(self, registry, raw): self.timestamps = [datetime.fromtimestamp(x // 1000) for x in raw['ts']] self.samples = {} parameters = ((x, raw['units'][x]) for x in raw if x not in self._INTERNAL_FIELDS) for parameter, unit in parameters: self.samples[parameter] = [x * registry(_convert_notation(unit)) for x in raw[parameter]] def __len__(self): return len(self.timestamps) def parameters(self) -> tuple: return tuple(self.samples.keys()) def entries(self) -> EntryView: return EntryView(self) def __iter__(self): return self.entries() @dataclass class Endpoint: path: str def __call__(self, ctx, *args, **kwargs): try: body = args[0].json() except (IndexError, AttributeError): body = Request(*args, **kwargs).json() response = ctx.session.post(ctx.api + self.path, json=body) response.raise_for_status() return Response(ctx.registry, response.json())