Added wrapper for responses

1 files changed, 55 insertions, 1 deletions

diff --git a/windy/point_forecast.py b/windy/point_forecast.py
index 7f1930a..8d55c5f 100644
--- a/windy/point_forecast.py
+++ b/windy/point_forecast.py
@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from datetime import datetime
 from enum import Enum
 
 
@@ -67,6 +68,57 @@ class Request:
 		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, registry(raw['units'][x])) for x in raw if x not in self._INTERNAL_FIELDS)
+		for parameter, unit in parameters:
+			self.samples[parameter] = [x * 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
@@ -76,4 +128,6 @@ class Endpoint:
 			body = args[0].json()
 		except (IndexError, AttributeError):
 			body = Request(*args, **kwargs).json()
-		return ctx.session.post(ctx.api + self.path, json=body)
+		response = ctx.session.post(ctx.api + self.path, json=body)
+		response.raise_for_status()
+		return Response(ctx.registry, response.json())