diff options
| -rw-r--r-- | pyproject.toml | 3 | ||||
| -rw-r--r-- | windy/point_forecast.py | 28 |
2 files changed, 22 insertions, 9 deletions
diff --git a/pyproject.toml b/pyproject.toml index a5df547..b7a809a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,8 @@ name="windy" dependencies=[ "requests==2.*", - "pint" + "pint==0.19.*", + "numpy==1.*", ] dynamic=["version"] diff --git a/windy/point_forecast.py b/windy/point_forecast.py index 0132751..308d432 100644 --- a/windy/point_forecast.py +++ b/windy/point_forecast.py @@ -6,6 +6,8 @@ from dataclasses import dataclass from datetime import datetime from enum import Enum +import numpy as np + def _json(value): try: @@ -82,6 +84,11 @@ class Level(_StrEnum): H200 = "200h" H150 = "150h" + def pressure(self): + if self is Level.SURFACE: + raise ValueError + return float(self.value[:-1]) + @dataclass class Request: @@ -132,7 +139,7 @@ class Prediction: return iter(self.parameters) def __getitem__(self, key): - return self._response.samples[key][self._index] + return self._response.values[key][self._index] class Response: @@ -143,7 +150,7 @@ class Response: Can be used in a for-loop to access all samples via Prediction: >>> for prediction in response: - >>> print(prediction.timestamp, prediction['temp-surface']) + >>> print(prediction.timestamp, prediction['temp']) Otherwise, timestamps list and samples dictionary are available for direct access. """ @@ -151,12 +158,17 @@ class Response: def __init__(self, registry, raw): self.timestamps = [datetime.fromtimestamp(x // 1000) for x in raw['ts']] - self.parameters = tuple(x for x in raw if x not in self._INTERNAL_FIELDS) - self.levels = tuple(sorted({Level(x.split("-")[1]) for x in self.parameters})) - self.samples = {} - for parameter in self.parameters: - unit = registry(_convert_notation(raw['units'][parameter])) - self.samples[parameter] = [x * unit for x in raw[parameter]] + count = len(self.timestamps) + split = [tuple(x.split("-")) for x in raw if x not in self._INTERNAL_FIELDS] + self.parameters = tuple({x for x, _ in split}) + self.levels = tuple(sorted({Level(x) for _, x in split})) + if 'pressure' in self.parameters: + for level in self.levels: + if level is Level.SURFACE: + continue + raw[f'pressure-{level}'] = [level.pressure() for _ in range(len(self.timestamps))] + units = {x: registry(_convert_notation(raw['units'][f'{x}-surface'])) for x in self.parameters} # Don't guess surface + self.values = {p: [[raw[f'{p}-{l}'][i] for l in self.levels] * units[p] for i in range(count)] for p in self.parameters} def __len__(self): return len(self.timestamps) |