diff --git a/telemetry_sh/telemetry.py b/telemetry_sh/telemetry.py index 497e509..4921436 100644 --- a/telemetry_sh/telemetry.py +++ b/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.") diff --git a/telemetry_sh/telemetry_async.py b/telemetry_sh/telemetry_async.py index bcbc458..e0e50cb 100644 --- a/telemetry_sh/telemetry_async.py +++ b/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." diff --git a/telemetry_sh/telemetry_types.py b/telemetry_sh/telemetry_types.py new file mode 100644 index 0000000..f53ef29 --- /dev/null +++ b/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 +