Type fixes for util modules (#322)

* fix typing for utils
* fix typing for main

Author: pgoellnitz

main.py

@@ -15,7 +15,7 @@
 from phenoback.utils import glogging
 
 
-def sentry_environment() -> tuple[str, float]:
+def sentry_environment() -> tuple[str, float, float]:
     project = g.get_project()
     if project == "phaenonet":
         return ("production", 1.0, 0.0)
@@ -38,11 +38,11 @@ def sentry_environment() -> tuple[str, float]:
     options={"storageBucket": os.environ.get("storageBucket")}
 )
 
-log: logging.Logger = None  # pylint: disable=invalid-name
+log: logging.Logger = None  # type: ignore # pylint: disable=invalid-name
 
 
 @contextmanager  # workaround as stackdriver fails to capture stackstraces
-def setup(data: dict | Request | None, context=None, level=logging.DEBUG):
+def setup(data: str | dict | Request | None, context=None, level=logging.DEBUG):
     """
     Setup logging and and capture exceptions.
     :param data: May be a dict, a http request or None.
@@ -56,8 +56,10 @@ def setup(data: dict | Request | None, context=None, level=logging.DEBUG):
             log.debug(context)
         if data:
             if isinstance(data, Request):
-                data = data.json if data.is_json else data.data
-            log.debug(data)
+                data_logged = data.json if data.is_json else data.data
+            else:
+                data_logged = data
+            log.debug(data_logged)
         yield
     except Exception:
         log.exception("Fatal error in cloud function")

phenoback/functions/iot/app.py

@@ -2,7 +2,7 @@
 import logging
 from zoneinfo import ZoneInfo
 
-import google
+import google.api_core.exceptions
 
 import phenoback.utils.data as d
 import phenoback.utils.firestore as f
@@ -27,7 +27,7 @@ def main_individual_updated(data, context):
         if g.get_field(data, "deveui", expected=False):
             individual = g.get_field(data, "individual")
             deveui = g.get_field(data, "deveui")
-            sensor_set(individual_id, individual, deveui)
+            sensor_set(individual_id, individual, deveui)  # type: ignore
         else:
             remove_sensor(individual_id)
 

phenoback/functions/iot/decoder.py

@@ -14,11 +14,11 @@ def get_value(self, start, length, signed=False):
         return value
 
     @property
-    def data(self):
+    def data(self) -> dict:
         return self._data
 
     @property
-    def uplink(self) -> dict:
+    def uplink(self) -> dict | None:
         return self.data.get("DevEUI_uplink")
 
     @property
@@ -30,7 +30,7 @@ def payload(self) -> str:
         return self.data.get("DevEUI_uplink", {}).get("payload_hex")
 
     @property
-    def decoded_payload(self) -> str:
+    def decoded_payload(self) -> dict:
         return self.data.get("DevEUI_uplink", {}).get("decoded_payload")
 
     def _set_decoded_payload(self, data: dict) -> None:

phenoback/functions/iot/dragino.py

@@ -18,10 +18,10 @@
 
 
 def main(request: Request):
-    if request.is_json:
+    if request.is_json and request.json:
         process_dragino(request.json)
     else:  # pragma: no cover
-        log.error("No json headers found")
+        log.error("No json headers set or payload is None")
         return Response("No json payload", HTTPStatus.BAD_REQUEST)
     return Response("ok", HTTPStatus.OK)
 
@@ -52,7 +52,7 @@ def process_dragino(data: dict) -> None:
         log.debug("No uplink data, skip")
 
 
-def set_uplink_frequency(deveui: str, interval: int, at: datetime = None):
+def set_uplink_frequency(deveui: str, interval: int, at: datetime | None = None):
     log.info("set uplink frequency to %is for %s at %s", interval, deveui, at)
     task_client().send(
         "",
@@ -61,7 +61,7 @@ def set_uplink_frequency(deveui: str, interval: int, at: datetime = None):
             "Payload": f"{16777216 + interval:{0}8x}",
             "FPort": 1,
         },
-        at=at,
+        at=at,  # type: ignore
     )
 
 

phenoback/utils/data.py

@@ -56,7 +56,7 @@ def get_phenoyear(reset_cache=False) -> int:
     return _get_dynamic_config()["phenoyear"]
 
 
-def update_phenoyear(year: int, transaction: Transaction = None) -> None:
+def update_phenoyear(year: int, transaction: Transaction | None = None) -> None:
     write_document(
         "definitions",
         "config_dynamic",
@@ -67,11 +67,15 @@ def update_phenoyear(year: int, transaction: Transaction = None) -> None:
     _get_dynamic_config.cache_clear()
 
 
-def get_individual(individual_id: str, transaction: Transaction = None) -> dict:
+def get_individual(
+    individual_id: str, transaction: Transaction | None = None
+) -> dict | None:
     return get_document("individuals", individual_id, transaction=transaction)
 
 
-def delete_individual(individual_id: str, transaction: Transaction = None) -> None:
+def delete_individual(
+    individual_id: str, transaction: Transaction | None = None
+) -> None:
     delete_document("individuals", individual_id, transaction=transaction)
 
 
@@ -84,7 +88,7 @@ def query_individuals(field_path: str, op_string: str, value: Any) -> Query:
 
 
 def write_individuals(
-    individuals: list[dict], key: str, transaction: Transaction = None
+    individuals: list[dict], key: str, transaction: Transaction | None = None
 ) -> None:
     if transaction is not None:
         write_batch("individuals", key, individuals, transaction=transaction)
@@ -93,35 +97,42 @@ def write_individuals(
 
 
 def write_individual(
-    individual_id: str, data: dict, merge: bool = False, transaction: Transaction = None
+    individual_id: str,
+    data: dict,
+    merge: bool = False,
+    transaction: Transaction | None = None,
 ) -> None:
     write_document(
         "individuals", individual_id, data, merge=merge, transaction=transaction
     )
 
 
 def update_individual(
-    individual_id: str, data: dict, transaction: Transaction = None
+    individual_id: str, data: dict, transaction: Transaction | None = None
 ) -> None:
     update_document("individuals", individual_id, data, transaction=transaction)
 
 
-def get_observation(observation_id: str, transaction: Transaction = None) -> dict:
+def get_observation(
+    observation_id: str, transaction: Transaction | None = None
+) -> dict | None:
     return get_document("observations", observation_id, transaction=transaction)
 
 
 def update_observation(
-    observation_id: str, data: dict, transaction: Transaction = None
+    observation_id: str, data: dict, transaction: Transaction | None = None
 ) -> None:
     update_document("observations", observation_id, data, transaction=transaction)
 
 
-def delete_observation(observation_id: str, transaction: Transaction = None) -> None:
+def delete_observation(
+    observation_id: str, transaction: Transaction | None = None
+) -> None:
     delete_document("observations", observation_id, transaction=transaction)
 
 
 def write_observation(
-    observation_id: str, data: dict, transaction: Transaction = None
+    observation_id: str | None, data: dict, transaction: Transaction | None = None
 ) -> None:
     write_document("observations", observation_id, data, transaction=transaction)
 
@@ -145,7 +156,7 @@ def create_user(
     )
 
 
-def get_user(user_id: str, transaction: Transaction = None) -> dict:
+def get_user(user_id: str, transaction: Transaction | None = None) -> dict | None:
     return get_document("users", user_id, transaction=transaction)
 
 
@@ -166,7 +177,7 @@ def get_user_id_by_email(email: str) -> str:  # pragma: no cover
 
 
 def follow_user(
-    follower_id: str, followee_id: str, transaction: Transaction = None
+    follower_id: str, followee_id: str, transaction: Transaction | None = None
 ) -> bool:
     user = get_user(follower_id, transaction=transaction)
     if not user:
@@ -192,17 +203,17 @@ def has_sensor(individual: dict) -> bool:
     return individual.get("sensor") is not None
 
 
-def localtime(timestamp: datetime = None):
+def localtime(timestamp: datetime | None = None):
     if not timestamp:
         timestamp = datetime.now(tz=tzlocal.get_localzone())
     return timestamp.astimezone(ZoneInfo("Europe/Zurich"))
 
 
-def localdate(timestamp: datetime = None):
+def localdate(timestamp: datetime | None = None):
     return localtime(timestamp).date()
 
 
-def to_id_array(data: dict[dict], key="id") -> list[dict]:
+def to_id_array(data: dict[str, dict], key="id") -> list[dict]:
     """
     Convert a dictionary to an array of dictionaries with an additional key.
     Useful for writing data in batch mode.

phenoback/utils/firestore.py

@@ -55,7 +55,7 @@ def transaction_commit():
 
 
 def delete_document(
-    collection: str, document_id: str, transaction: Transaction = None
+    collection: str, document_id: str, transaction: Transaction | None = None
 ) -> None:
     log.debug("Delete document %s from %s", document_id, collection)
     ref = firestore_client().collection(collection).document(document_id)
@@ -80,8 +80,8 @@ def _delete_batch(coll_ref, batch_size: int = 1000):
         return None
 
 
-def delete_collection(collection: str, batch_size: int = 1000) -> None:
-    _delete_batch(get_collection(collection), batch_size)
+def delete_collection(collection_name: str, batch_size: int = 1000) -> None:
+    _delete_batch(collection(collection_name), batch_size)
 
 
 def delete_batch(
@@ -125,8 +125,8 @@ def write_batch(
     data: list[dict],
     *,
     merge: bool = False,
-    commit_size: int = None,
-    transaction: Transaction = None,
+    commit_size: int | None = None,
+    transaction: Transaction | None = None,
     commit_sleep: float = 0,
 ) -> int:
     if transaction is not None:
@@ -168,7 +168,7 @@ def write_document(
     document_id: str | None,
     data: dict,
     merge: bool = False,
-    transaction: Transaction = None,
+    transaction: Transaction | None = None,
 ) -> None:
     log.debug(
         "Write document %s to %s (%s)",
@@ -187,7 +187,7 @@ def update_document(
     collection: str,
     document_id: str,
     data: dict,
-    transaction: Transaction = None,
+    transaction: Transaction | None = None,
 ) -> None:
     log.debug(
         "Update document %s in %s (%s)",
@@ -203,7 +203,7 @@ def update_document(
 
 
 def get_document(
-    collection: str, document_id: str, transaction: Transaction = None
+    collection: str, document_id: str, transaction: Transaction | None = None
 ) -> dict | None:
     log.debug(
         "Get document %s in %s (%s)",
@@ -220,7 +220,7 @@ def get_document(
     )
 
 
-def collection(collection: str) -> Query:
+def collection(collection: str) -> CollectionReference:
     log.debug("Query %s", collection)
     return firestore_client().collection(collection)
 
@@ -236,13 +236,8 @@ def query_collection(
     )
 
 
-def get_collection(collection: str) -> CollectionReference:
-    log.debug("Query collection %s", collection)
-    return firestore_client().collection(collection)
-
-
-def get_collection_documents(collection: str) -> list[dict]:
-    return [location.to_dict() for location in get_collection(collection).stream()]
+def get_collection_documents(collection_name: str) -> list[dict]:
+    return [location.to_dict() for location in collection(collection_name).stream()]
 
 
 def docs2str(docs):  # pragma: no cover

phenoback/utils/gcloud.py

@@ -3,7 +3,6 @@
 import json
 import logging
 import os
-from collections import namedtuple
 from datetime import datetime
 
 import dateparser
@@ -15,7 +14,7 @@
 
 def get_field(
     data: dict, fieldname: str, old_value: bool = False, expected=True
-) -> str | int | datetime | None:
+) -> str | int | float | datetime | bool | dict | list | None:
     value_type = "oldValue" if old_value else "value"
     value_dict = data[value_type].get("fields", {}).get(fieldname)
     if value_dict is not None:
@@ -28,7 +27,9 @@ def get_field(
         return None
 
 
-def _get_field(value_dict: dict):
+def _get_field(
+    value_dict: dict,
+) -> str | int | float | datetime | bool | dict | list | None:
     value = next(iter(value_dict.values()))
     value_type = next(iter(value_dict.keys()))
     if value_type == "stringValue":
@@ -63,7 +64,7 @@ def context2dict(context: Context) -> dict:
     return context.__dict__
 
 
-def dict2context(context_dict) -> namedtuple:
+def dict2context(context_dict) -> Context:
     return Context(
         eventId=context_dict.get("event_id"),
         timestamp=context_dict.get("timestamp"),

phenoback/utils/pubsub.py

@@ -17,9 +17,7 @@ def __init__(self, topic: str) -> None:
         self.topic_path = self.client.topic_path(self.project, self.topic)
         log.debug("Publisher created for topic %s", self.topic)
 
-    def send(
-        self, payload: dict | str, metadata: dict[str, bytes | str] = None
-    ) -> None:
+    def send(self, payload: dict | str, metadata: dict[str, str] | None = None) -> None:
         if metadata is None:
             metadata = {}
         if isinstance(payload, dict):