Make access token awaitable

Author: Bre77

teslemetry_stream/stream.py

@@ -1,16 +1,17 @@
-from collections.abc import Callable
-from typing import Any
-import aiohttp
 import asyncio
 import json
 import logging
 from datetime import datetime, timezone
+from typing import Any, Awaitable, Callable
+
+import aiohttp
 
-from .vehicle import TeslemetryStreamVehicle
 from .exception import TeslemetryStreamEnded
+from .vehicle import TeslemetryStreamVehicle
 
 LOGGER = logging.getLogger(__package__)
 
+
 class TeslemetryStream:
     """Teslemetry Stream Client"""
 
@@ -19,7 +20,7 @@ class TeslemetryStream:
     def __init__(
         self,
         session: aiohttp.ClientSession,
-        access_token: str,
+        access_token: str | Callable[[], Awaitable[str]],
         server: str = "api.teslemetry.com",
         vin: str | None = None,
         parse_timestamp: bool = False,
@@ -41,19 +42,30 @@ def __init__(
         self.active: bool = False
         self.server = server
         self.vin = vin
-        self._listeners: dict[Callable, tuple[Callable[[dict[str,Any]],None], dict | None]] = {}
-        self._connection_listeners: dict[Callable, Callable[[bool],None]] = {}
+        self._listeners: dict[
+            Callable, tuple[Callable[[dict[str, Any]], None], dict | None]
+        ] = {}
+        self._connection_listeners: dict[Callable, Callable[[bool], None]] = {}
         self._session = session
-        self._headers = {"Authorization": f"Bearer {access_token}", "X-Library": "python teslemetry-stream"}
+        self.access_token = access_token
         self.parse_timestamp = parse_timestamp
         self.manual = manual
         self.retries: int = 0
         self.vehicles: dict[str, TeslemetryStreamVehicle] = {}
 
-        if(self.vin):
+        if self.vin:
             self.vehicle: TeslemetryStreamVehicle = self.get_vehicle(self.vin)
             self.vehicles[self.vin] = self.vehicle
 
+    async def headers(self) -> dict[str, str]:
+        if callable(self.access_token):
+            access_token = await self.access_token()
+        else:
+            access_token = self.access_token
+        return {
+            "Authorization": f"Bearer {access_token}",
+            "X-Library": "python teslemetry-stream",
+        }
 
     def get_vehicle(self, vin: str) -> TeslemetryStreamVehicle:
         """
@@ -83,20 +95,21 @@ async def get_config(self, vin: str | None = None) -> None:
         """
         if not self.server:
             await self.find_server()
-        if hasattr(self, 'vehicle'):
+        if hasattr(self, "vehicle"):
             await self.vehicle.get_config()
 
     async def find_server(self) -> None:
         """
         Find the server using metadata.
         """
+        headers = await self.headers()
         req = await self._session.get(
             "https://api.teslemetry.com/api/metadata",
-            headers=self._headers,
+            headers=headers,
             raise_for_status=True,
         )
         response = await req.json()
-        self.server = f"{response["region"].lower()}.teslemetry.com"
+        self.server = f"{response['region'].lower()}.teslemetry.com"
 
     async def update_fields(self, fields: dict, vin: str) -> dict:
         """
@@ -106,9 +119,10 @@ async def update_fields(self, fields: dict, vin: str) -> dict:
         :param vin: Vehicle Identification Number.
         :return: Response JSON as a dictionary.
         """
+        headers = await self.headers()
         resp = await self._session.patch(
             f"https://api.teslemetry.com/api/config/{self.vin}",
-            headers=self._headers,
+            headers=headers,
             json={"fields": fields},
             raise_for_status=False,
         )
@@ -124,9 +138,10 @@ async def replace_fields(self, fields: dict, vin: str) -> dict:
         :param vin: Vehicle Identification Number.
         :return: Response JSON as a dictionary.
         """
+        headers = await self.headers()
         resp = await self._session.post(
             f"https://api.teslemetry.com/api/config/{self.vin}",
-            headers=self._headers,
+            headers=headers,
             json={"fields": fields},
             raise_for_status=False,
         )
@@ -154,6 +169,7 @@ def async_add_connection_listener(
         :param callback: Callback function to handle connection state changes.
         :return: Function to remove the listener.
         """
+
         def remove_listener() -> None:
             """
             Remove connection listener.
@@ -181,14 +197,15 @@ async def connect(self) -> None:
         url = f"https://{self.server}/sse"
         if self.vin:
             url += f"/{self.vin}"
+        headers = await self.headers()
         self._response = await self._session.get(
             url,
-            headers=self._headers,
+            headers=headers,
             raise_for_status=True,
             timeout=aiohttp.ClientTimeout(
                 connect=5, sock_connect=5, sock_read=30, total=None
             ),
-            chunked=True
+            chunked=True,
         )
         LOGGER.debug(
             "Connected to %s with status %s", self._response.url, self._response.status
@@ -223,7 +240,7 @@ def __aiter__(self):
         self.active = True
         return self
 
-    async def __anext__(self) -> dict:
+    async def __anext__(self) -> dict[str, Any]:
         """
         Return next event.
 
@@ -262,7 +279,7 @@ async def __anext__(self) -> dict:
         except aiohttp.ClientError as error:
             LOGGER.warning("Client error: %s", repr(error))
             self.close()
-            delay = min(2 ** self.retries, 600)
+            delay = min(2**self.retries, 600)
             LOGGER.debug("Reconnecting in %s seconds", delay)
             await asyncio.sleep(delay)
             self.retries += 1
@@ -272,7 +289,6 @@ async def __anext__(self) -> dict:
             LOGGER.debug("Reconnecting in %s seconds", 1)
             await asyncio.sleep(1)
 
-
     def async_add_listener(
         self, callback: Callable, filters: dict | None = None
     ) -> Callable[[], None]:
@@ -316,16 +332,17 @@ async def listen(self):
                             LOGGER.error("Uncaught error in listener: %s", error)
         LOGGER.debug("Listen has finished")
 
-    def listen_Credits(self, callback: Callable[[dict[str, str | int]], None]) -> Callable[[], None]:
+    def listen_Credits(
+        self, callback: Callable[[dict[str, str | int]], None]
+    ) -> Callable[[], None]:
         """
         Listen for credits update.
 
         :param callback: Callback function to handle credits update.
         :return: Function to remove the listener.
         """
         return self.async_add_listener(
-            lambda x: callback(x["credits"]),
-            {"credits": None}
+            lambda x: callback(x["credits"]), {"credits": None}
         )
 
     def listen_Balance(self, callback: Callable[[int], None]) -> Callable[[], None]:
@@ -336,10 +353,10 @@ def listen_Balance(self, callback: Callable[[int], None]) -> Callable[[], None]:
         :return: Function to remove the listener.
         """
         return self.async_add_listener(
-            lambda x: callback(x["credits"]["balance"]),
-            {"credits": {"balance": None}}
+            lambda x: callback(x["credits"]["balance"]), {"credits": {"balance": None}}
         )
 
+
 def recursive_match(dict1, dict2):
     """
     Recursively match dict1 with dict2.

teslemetry_stream/vehicle.py

@@ -84,9 +84,10 @@ def config(self) -> dict:
     async def get_config(self) -> None:
         """Get the current configuration for the vehicle."""
 
+        headers = await self.stream.headers()
         req = await self.stream._session.get(
             f"https://api.teslemetry.com/api/config/{self.vin}",
-            headers=self.stream._headers,
+            headers=headers,
             raise_for_status=False,
         )
         if req.status == 200:
@@ -133,19 +134,21 @@ async def update_config(self, config: dict) -> None:
 
     async def patch_config(self, config: dict) -> dict[str, str | dict]:
         """Modify the configuration for the vehicle."""
+        headers = await self.stream.headers()
         resp = await self.stream._session.patch(
             f"https://api.teslemetry.com/api/config/{self.vin}",
-            headers=self.stream._headers,
+            headers=headers,
             json=config,
             raise_for_status=False,
         )
         return await resp.json()
 
     async def post_config(self, config: dict) -> dict[str, str | dict]:
         """Overwrite the configuration for the vehicle."""
+        headers = await self.stream.headers()
         resp = await self.stream._session.post(
             f"https://api.teslemetry.com/api/config/{self.vin}",
-            headers=self.stream._headers,
+            headers=headers,
             json=config,
             raise_for_status=False,
         )