Add async context for Home Assistant

README.md

@@ -41,17 +41,15 @@ from bcontrolpy import BControl, AuthenticationError
 
 async def main():
     # Connect to EM300 meter on local network
-    bc = BControl(ip="192.168.1.100", password="your_password")
-    try:
-        info = await bc.login()
-        print("Login successful:", info)
-
-        data = await bc.get_data()
-        print("Meter readings:", data)
-    except AuthenticationError:
-        print("Authentication failed: check your credentials")
-    finally:
-        await bc.close()
+    async with BControl(ip="192.168.1.100", password="your_password") as bc:
+        try:
+            info = await bc.login()
+            print("Login successful:", info)
+
+            data = await bc.get_data()
+            print("Meter readings:", data)
+        except AuthenticationError:
+            print("Authentication failed: check your credentials")
 
 asyncio.run(main())
 ```
@@ -69,10 +67,9 @@ asyncio.run(main())
 Example:
 
 ```python
-bc = BControl(ip="192.168.1.100", password="your_password")
-info = await bc.login()
-values = await bc.get_data()
-await bc.close()
+async with BControl(ip="192.168.1.100", password="your_password") as bc:
+    info = await bc.login()
+    values = await bc.get_data()
 ```
 
 ### Command Line Example

example/example.py

@@ -5,11 +5,11 @@
 import json  # Add this import at the top
 
 
-async def main(ip, password):    
+async def main(ip, password):
     async with aiohttp.ClientSession() as session:
-        bc = BControl(ip, password, session=session)
-        login_response = await bc.login()
-        print("Login Response:", login_response)
+        async with BControl(ip, password, session=session) as bc:
+            login_response = await bc.login()
+            print("Login Response:", login_response)
 
         try:
             while True:
@@ -19,8 +19,6 @@ async def main(ip, password):
                 await asyncio.sleep(5)  # Wait for 5 seconds before the next call
         except asyncio.CancelledError:
             print("Task cancelled, closing connection.")
-        finally:
-            await bc.close()
 
 
 def format_data(data):

pyproject.toml

@@ -14,7 +14,8 @@ readme = "README.md"
 requires-python = ">=3.8"
 
 dependencies = [
-    "aiohttp>=3.8"
+    "aiohttp>=3.8",
+    "async-timeout>=4.0"
 ]
 
 [tool.setuptools.packages.find]

src/bcontrolpy/__init__.py

@@ -1,3 +1,3 @@
 from .bcontrolpy import BControl
 
-__all__ = ["BControl"]
+__all__ = ["BControl"]

src/bcontrolpy/bcontrolpy.py

@@ -2,6 +2,7 @@
 import aiohttp
 import json
 import logging
+from async_timeout import timeout
 from .key_mapping import key_mapping
 
 _LOGGER = logging.getLogger(__name__)
@@ -24,10 +25,10 @@ class NotAuthenticatedError(Exception):
     """Raised when trying to get data without authentication."""
     pass
 
-async def getcookie(base_url: str):
+async def getcookie(session: aiohttp.ClientSession, base_url: str, timeout_seconds: int):
     url = f"{base_url}/start.php"
     try:
-        async with aiohttp.ClientSession() as session:
+        async with timeout(timeout_seconds):
             async with session.get(url) as resp:
                 resp.raise_for_status()
                 return resp.cookies, await resp.text()
@@ -38,17 +39,18 @@ async def getcookie(base_url: str):
     except Exception as e:
         raise CookieRetrievalError(f"Unexpected error during cookie retrieval: {e}")
 
-async def authenticate(session: aiohttp.ClientSession, base_url: str, login: str, password: str, cookie_value: str):
+async def authenticate(session: aiohttp.ClientSession, base_url: str, login: str, password: str, cookie_value: str, timeout_seconds: int):
     url = f"{base_url}/start.php"
     headers = {'Content-Type': 'application/x-www-form-urlencoded', 'Cookie': f'PHPSESSID={cookie_value}'}
     data = {'login': login, 'password': password}
     try:
-        async with session.post(url, data=data, headers=headers) as resp:
-            # Spezielles Handling für falsche Anmeldedaten
-            if resp.status == 403:
-                raise AuthenticationError("Invalid credentials: access forbidden (403)")
-            resp.raise_for_status()
-            return await resp.text()
+        async with timeout(timeout_seconds):
+            async with session.post(url, data=data, headers=headers) as resp:
+                # Spezielles Handling für falsche Anmeldedaten
+                if resp.status == 403:
+                    raise AuthenticationError("Invalid credentials: access forbidden (403)")
+                resp.raise_for_status()
+                return await resp.text()
     except AuthenticationError:
         raise
     except aiohttp.ClientResponseError as e:
@@ -60,52 +62,42 @@ async def authenticate(session: aiohttp.ClientSession, base_url: str, login: str
     except Exception as e:
         raise AuthenticationError(f"Unexpected error during authentication: {e}")
 
-async def getdata(session: aiohttp.ClientSession, base_url: str, cookie_value: str):
+async def getdata(session: aiohttp.ClientSession, base_url: str, cookie_value: str, timeout_seconds: int):
     url = f"{base_url}/mum-webservice/data.php"
     headers = {'Cookie': f'PHPSESSID={cookie_value}'}
-    async with session.get(url, headers=headers) as resp:
-        resp.raise_for_status()
-        return await resp.text()
+    async with timeout(timeout_seconds):
+        async with session.get(url, headers=headers) as resp:
+            resp.raise_for_status()
+            return await resp.text()
 
 
 def translate_keys(data: dict, mapping: dict) -> dict:
     return {mapping.get(k, k): v for k, v in data.items()}
 
 class BControl:
-    def __init__(self, ip: str, password: str, session: aiohttp.ClientSession = None):
-        """Initialize the client.
-
-        Parameters
-        ----------
-        ip: str
-            IP address of the B-Control meter.
-        password: str
-            Password used for authentication.
-        session: aiohttp.ClientSession, optional
-            Reuse an existing session. If omitted, ``BControl`` creates its
-            own ``ClientSession`` which will be closed on :meth:`close`.
-        """
-
+    def __init__(self, ip: str, password: str, session: aiohttp.ClientSession | None = None, timeout_seconds: int = 10):
         self.base_url = f"http://{ip}"
         self.password = password
-        if session is None:
-            self.session = aiohttp.ClientSession()
-            self._session_owner = True
-        else:
-            self.session = session
-            self._session_owner = False
+        self.session = session or aiohttp.ClientSession()
+        self.timeout = timeout_seconds
         self.cookie_value = None
         self.logged_in = False
         self.serial = None
         self.app_version = None
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.close()
+
     async def login(self) -> dict:
         """
         Logs in and returns a dict with serial, app_version and authentication status.
         Raises AuthenticationError if credentials are invalid.
         """
         try:
-            cookies, text = await getcookie(self.base_url)
+            cookies, text = await getcookie(self.session, self.base_url, self.timeout)
             init_data = json.loads(text)
             login_val = init_data.get("serial")
             if not login_val:
@@ -116,7 +108,7 @@ async def login(self) -> dict:
                 raise CookieValueError("PHPSESSID cookie missing after start.")
             self.cookie_value = phpsess.value
 
-            auth_text = await authenticate(self.session, self.base_url, login_val, self.password, self.cookie_value)
+            auth_text = await authenticate(self.session, self.base_url, login_val, self.password, self.cookie_value, self.timeout)
             auth = json.loads(auth_text)
 
             # nur die benötigten Felder
@@ -142,12 +134,12 @@ async def get_data(self) -> dict:
             _LOGGER.info("Session not valid, logging in first...")
             await self.login()
 
-        raw = await getdata(self.session, self.base_url, self.cookie_value)
+        raw = await getdata(self.session, self.base_url, self.cookie_value, self.timeout)
         data = json.loads(raw)
         if data.get("authentication") is False:
             _LOGGER.warning("Session expired, re-login")
             await self.login()
-            raw = await getdata(self.session, self.base_url, self.cookie_value)
+            raw = await getdata(self.session, self.base_url, self.cookie_value, self.timeout)
             data = json.loads(raw)
 
         return translate_keys(data, key_mapping)