add typing hints for telemetry-python

telemetry_sh/telemetry.py

@@ -1,6 +1,7 @@
 import requests
 import json
 from typing import Union, List
+from telemetry_sh.telemetry_types import TelemetryQueryResponse, TelemetryLogResponse, TelemetryQueryType
 
 class Telemetry:
     def __init__(self):
@@ -16,7 +17,7 @@ def check_and_return(self, response: requests.Response) -> dict:
             raise Exception(response_json.get("message", "Unknown error"))
         return response_json
 
-    def log(self, table: str, data: Union[dict, List[dict]]) -> dict:
+    def log(self, table: str, data: Union[dict, List[dict]]) -> TelemetryLogResponse:
         if not self.api_key:
             raise ValueError("API key is not initialized. Please call init() with your API key.")
 
@@ -33,7 +34,7 @@ def log(self, table: str, data: Union[dict, List[dict]]) -> dict:
         response = requests.post(f"{self.base_url}/log", headers=headers, data=json.dumps(body))
         return self.check_and_return(response)
 
-    def query(self, query: str) -> dict:
+    def query(self, query: str) -> TelemetryQueryResponse[TelemetryQueryType]:
         if not self.api_key:
             raise ValueError("API key is not initialized. Please call init() with your API key.")
 

telemetry_sh/telemetry_async.py

@@ -1,5 +1,6 @@
 import aiohttp
 from typing import Union, List
+from telemetry_sh.telemetry_types import TelemetryQueryResponse, TelemetryQueryType, TelemetryLogResponse
 
 
 class TelemetryAsync:
@@ -16,7 +17,7 @@ async def check_and_return(self, response: aiohttp.ClientResponse) -> dict:
             raise Exception(respose_json.get("message", "Unknown error"))
         return respose_json
 
-    async def log(self, table: str, data: Union[dict, List[dict]]) -> dict:
+    async def log(self, table: str, data: Union[dict, List[dict]]) -> TelemetryLogResponse:
         if not self.api_key:
             raise ValueError(
                 "API key is not initialized. Please call init() with your API key."
@@ -32,7 +33,7 @@ async def log(self, table: str, data: Union[dict, List[dict]]) -> dict:
             ) as response:
                 return await self.check_and_return(response)
 
-    async def query(self, query: str) -> dict:
+    async def query(self, query: str) -> TelemetryQueryResponse[TelemetryQueryType]:
         if not self.api_key:
             raise ValueError(
                 "API key is not initialized. Please call init() with your API key."

telemetry_sh/telemetry_types.py

@@ -0,0 +1,12 @@
+from typing import Mapping, TypedDict, Literal, TypeVar, Generic
+TelemetryQueryType = TypeVar("TelemetryQueryType", bound=Mapping)
+
+class TelemetryQueryResponse(TypedDict, Generic[TelemetryQueryType]):
+    data: list[TelemetryQueryType]
+    status: Literal["success"]
+    key_order: list[str]
+
+class TelemetryLogResponse(TypedDict):
+    status: Literal["success"]
+    message: str
+