refactor: add type hints and pydantic validation schemas

Author: lupppig

openpaygo/metrics_request.py

@@ -1,50 +1,66 @@
 import copy
-
+from typing import Optional, List, Dict, Any, Union
+from .models import MetricsDataFormat, MetricsHistoricalDataStep
 from .metrics_shared import OpenPAYGOMetricsShared
 
 
 class MetricsRequestHandler(object):
     def __init__(
-        self, serial_number, data_format=None, secret_key=None, auth_method=None
-    ):
+        self, serial_number: str, data_format: Optional[Union[Dict[str, Any], MetricsDataFormat]] = None, secret_key: Optional[str] = None, auth_method: Optional[str] = None
+    ) -> None:
         self.secret_key = secret_key
         self.auth_method = auth_method
-        self.request_dict = {
+        self.request_dict: Dict[str, Any] = {
             "serial_number": serial_number,
         }
-        self.data_format = data_format
+        if data_format is not None:
+            if isinstance(data_format, dict):
+                self.data_format: Optional[Dict[str, Any]] = MetricsDataFormat.model_validate(data_format).model_dump(exclude_none=True)
+            else:
+                self.data_format: Optional[Dict[str, Any]] = data_format.model_dump(exclude_none=True)
+        else:
+            self.data_format: Optional[Dict[str, Any]] = None
+            
         if self.data_format:
             if self.data_format.get("id"):
-                self.request_dict["data_format_id"] = data_format.get("id")
+                self.request_dict["data_format_id"] = self.data_format.get("id")
             else:
-                self.request_dict["data_format"] = data_format
-        self.data = {}
-        self.historical_data = {}
+                self.request_dict["data_format"] = self.data_format
+        self.data: Dict[str, Any] = {}
+        self.historical_data: List[Dict[str, Any]] = []
 
-    def set_request_count(self, request_count):
+    def set_request_count(self, request_count: int) -> None:
         self.request_dict["request_count"] = request_count
 
-    def set_timestamp(self, timestamp):
+    def set_timestamp(self, timestamp: int) -> None:
         self.request_dict["timestamp"] = timestamp
 
-    def set_data(self, data):
+    def set_data(self, data: Dict[str, Any]) -> None:
         self.data = data
 
-    def set_historical_data(self, historical_data):
-        if not self.data_format.get("historical_data_interval"):
-            for time_step in historical_data:
+    def set_historical_data(self, historical_data: List[Union[Dict[str, Any], MetricsHistoricalDataStep]]) -> None:
+        validated_historical_data = []
+        for time_step in historical_data:
+            if isinstance(time_step, dict):
+                step = MetricsHistoricalDataStep.model_validate(time_step).model_dump(exclude_none=True, mode='json')
+            else:
+                step = time_step.model_dump(exclude_none=True, mode='json')
+            validated_historical_data.append(step)
+            
+        if self.data_format and not self.data_format.get("historical_data_interval"):
+            for time_step in validated_historical_data:
                 if not time_step.get("timestamp"):
                     raise ValueError(
                         "Historical Data objects must have a time stamp if no "
                         "historical_data_interval is defined."
                     )
-        self.historical_data = historical_data
+        self.historical_data = validated_historical_data
 
-    def get_simple_request_payload(self):
+    def get_simple_request_payload(self) -> str:
         payload = self.get_simple_request_dict()
         return OpenPAYGOMetricsShared.convert_to_metrics_json(payload)
 
-    def get_simple_request_dict(self):
+    def get_simple_request_dict(self) -> Dict[str, Any]:
         simple_request = self.request_dict
         simple_request["data"] = self.data
         simple_request["historical_data"] = self.historical_data
@@ -57,11 +73,11 @@ def get_simple_request_dict(self):
             )
         return simple_request
 
-    def get_condensed_request_payload(self):
+    def get_condensed_request_payload(self) -> str:
         payload = self.get_condensed_request_dict()
         return OpenPAYGOMetricsShared.convert_to_metrics_json(payload)
 
-    def get_condensed_request_dict(self):
+    def get_condensed_request_dict(self) -> Dict[str, Any]:
         if not self.data_format:
             raise ValueError("No Data Format provided for condensed request")
         data_order = self.data_format.get("data_order")

openpaygo/metrics_response.py

@@ -1,19 +1,21 @@
 import copy
 import json
 from datetime import datetime, timedelta
+from typing import Optional, List, Dict, Any, Union
 
+from .models import MetricsDataFormat
 from .metrics_shared import OpenPAYGOMetricsShared
 
 
 class MetricsResponseHandler(object):
     def __init__(
         self,
-        received_metrics,
-        data_format=None,
-        secret_key=None,
-        last_request_count=None,
-        last_request_timestamp=None,
-    ):
+        received_metrics: str,
+        data_format: Optional[Union[Dict[str, Any], MetricsDataFormat]] = None,
+        secret_key: Optional[str] = None,
+        last_request_count: Optional[int] = None,
+        last_request_timestamp: Optional[int] = None,
+    ) -> None:
         self.received_metrics = received_metrics
         self.request_dict = json.loads(received_metrics)
         # We convert the base variable names to simple
@@ -28,38 +30,53 @@ def __init__(
             self.timestamp = self.request_dict.get("timestamp")
         self.response_dict = {}
         self.secret_key = secret_key
-        self.data_format = data_format
         self.last_request_count = last_request_count
         self.last_request_timestamp = last_request_timestamp
+
+        if data_format is not None:
+            if isinstance(data_format, dict):
+                self.data_format: Optional[Dict[str, Any]] = MetricsDataFormat.model_validate(data_format).model_dump(exclude_none=True)
+            else:
+                self.data_format: Optional[Dict[str, Any]] = data_format.model_dump(exclude_none=True)
+        else:
+            self.data_format: Optional[Dict[str, Any]] = None
+
         if not self.data_format and self.request_dict.get("data_format"):
-            self.data_format = self.request_dict.get("data_format")
+            df = self.request_dict.get("data_format")
+            if isinstance(df, dict):
+                self.data_format = MetricsDataFormat.model_validate(df).model_dump(exclude_none=True)
+            else:
+                self.data_format = df
 
-    def get_device_serial(self):
+    def get_device_serial(self) -> str:
         return self.request_dict.get("serial_number")
 
-    def get_data_format_id(self):
+    def get_data_format_id(self) -> Optional[int]:
         return self.request_dict.get("data_format_id")
 
-    def data_format_available(self):
+    def data_format_available(self) -> bool:
         return self.data_format is not None
 
     def set_device_parameters(
         self,
-        secret_key=None,
-        data_format=None,
-        last_request_count=None,
-        last_request_timestamp=None,
-    ):
+        secret_key: Optional[str] = None,
+        data_format: Optional[Union[Dict[str, Any], MetricsDataFormat]] = None,
+        last_request_count: Optional[int] = None,
+        last_request_timestamp: Optional[int] = None,
+    ) -> None:
         if secret_key:
             self.secret_key = secret_key
-        if data_format:
-            self.data_format = data_format
+        if data_format is not None:
+            if isinstance(data_format, dict):
+                self.data_format = MetricsDataFormat.model_validate(data_format).model_dump(exclude_none=True)
+            else:
+                self.data_format = data_format.model_dump(exclude_none=True)
         if last_request_count:
             self.last_request_count = last_request_count
         if last_request_timestamp:
             self.last_request_timestamp = last_request_timestamp
 
-    def is_auth_valid(self):
+    def is_auth_valid(self) -> bool:
         auth_string = self.request_dict.get("auth", None)
         if not auth_string:
             return False
@@ -89,7 +106,7 @@ def is_auth_valid(self):
                 return True
         return False
 
-    def get_simple_metrics(self):
+    def get_simple_metrics(self) -> Dict[str, Any]:
         # We start the process by making a copy of the dict to work with
         simple_dict = copy.deepcopy(self.request_dict)
         simple_dict.pop("auth") if "auth" in simple_dict else None  # We remove the auth
@@ -103,34 +120,34 @@ def get_simple_metrics(self):
         )
         return simple_dict
 
-    def get_data_timestamp(self):
+    def get_data_timestamp(self) -> int:
         return self.request_dict.get("data_collection_timestamp", self.timestamp)
 
-    def get_request_timestamp(self):
+    def get_request_timestamp(self) -> Optional[int]:
         return self.request_timestamp
 
-    def get_request_count(self):
+    def get_request_count(self) -> Optional[int]:
         return self.request_dict.get("request_count")
 
-    def get_token_count(self):
+    def get_token_count(self) -> Optional[int]:
         data = self._get_simple_data()
         return data.get("token_count")
 
-    def expects_token_answer(self):
+    def expects_token_answer(self) -> bool:
         return self.get_token_count() is not None
 
-    def add_tokens_to_answer(self, token_list):
+    def add_tokens_to_answer(self, token_list: List[str]) -> None:
         self.response_dict["token_list"] = token_list
 
-    def expects_time_answer(self):
+    def expects_time_answer(self) -> bool:
         data = self._get_simple_data()
         if data.get("active_until_timestamp_requested", False) or data.get(
             "active_seconds_left_requested", False
         ):
             return True
         return False
 
-    def add_time_to_answer(self, target_datetime):
+    def add_time_to_answer(self, target_datetime: datetime) -> None:
         data = self._get_simple_data()
         if data.get("active_until_timestamp_requested", False):
             target_timestamp = 0
@@ -150,24 +167,24 @@ def add_time_to_answer(self, target_datetime):
         else:
             raise ValueError("No time requested")
 
-    def add_new_base_url_to_answer(self, new_base_url):
+    def add_new_base_url_to_answer(self, new_base_url: str) -> None:
         self.add_settings_to_answer({"base_url": new_base_url})
 
-    def add_settings_to_answer(self, settings_dict):
+    def add_settings_to_answer(self, settings_dict: Dict[str, Any]) -> None:
         if not self.response_dict.get("settings"):
             self.response_dict["settings"] = {}
         self.response_dict["settings"].update(settings_dict)
 
-    def add_extra_data_to_answer(self, extra_data_dict):
+    def add_extra_data_to_answer(self, extra_data_dict: Dict[str, Any]) -> None:
         if not self.response_dict.get("extra_data"):
             self.response_dict["extra_data"] = {}
         self.response_dict["extra_data"].update(extra_data_dict)
 
-    def get_answer_payload(self):
+    def get_answer_payload(self) -> str:
         payload = self.get_answer_dict()
         return OpenPAYGOMetricsShared.convert_to_metrics_json(payload)
 
-    def get_answer_dict(self):
+    def get_answer_dict(self) -> Dict[str, Any]:
         # If there is not data format, we just return the full response
         condensed_answer = copy.deepcopy(self.response_dict)
         if self.secret_key:

openpaygo/models.py

@@ -0,0 +1,27 @@
+from typing import Optional, List, Any, Dict
+from pydantic import BaseModel, ConfigDict
+
+
+class MetricsDataFormat(BaseModel):
+    id: Optional[int] = None
+    data_order: Optional[List[str]] = None
+    historical_data_order: Optional[List[str]] = None
+    historical_data_interval: Optional[int] = None
+
+
+class MetricsHistoricalDataStep(BaseModel):
+    timestamp: Optional[int] = None
+    relative_time: Optional[int] = None
+    model_config = ConfigDict(extra='allow')
+
+
+class MetricsRequestData(BaseModel):
+    serial_number: str
+    data_format_id: Optional[int] = None
+    data_format: Optional[MetricsDataFormat] = None
+    data: Optional[Dict[str, Any]] = None
+    historical_data: Optional[List[Any]] = None
+    request_count: Optional[int] = None
+    timestamp: Optional[int] = None
+    auth: Optional[str] = None
+    model_config = ConfigDict(extra='allow')