Accept existing aiohttp ClientSession. Implement api timeout.

pysabnzbd/__init__.py

@@ -1,67 +1,100 @@
 import aiohttp
+import asyncio
 
+from async_timeout import timeout
 
-class SabnzbdApi(object):
-    def __init__(self, base_url, api_key, web_root=''):
-        if base_url.endswith('/'):
-            base_url = base_url[:-1]
 
-        web_root = web_root.strip('/') + '/'
+class SabnzbdApi(object):
+    def __init__(self, base_url, api_key, web_root=None,
+                 session=None, timeout=5):
+        """Initialize the connection to the SABnzbd usenet client"""
+        if web_root is not None:
+            web_root = '{}/'.format(web_root.strip('/'))
+        else:
+            web_root = ''
 
-        self.api_url = '{}/{}{}'.format(base_url, web_root, 'api')
         self.queue = {}
-        self.default_params = {'apikey': api_key, 'output': 'json'}
+        self._api_url = '{}/{}{}'.format(base_url.rstrip('/'), web_root, 'api')
+        self._default_params = {'apikey': api_key, 'output': 'json'}
+        self._timeout = timeout
+
+        if session is None:
+            loop = asyncio.get_event_loop()
+            self._session = aiohttp.ClientSession(loop=loop)
+            self._cleanup_session = True
+        else:
+            self._session = session
+            self._cleanup_session = False
+
+    def __del__(self):
+        """Cleanup the session if it was created here"""
+        if self._cleanup_session:
+            self._session.loop.run_until_complete(self._session.close())
+
+    async def _call(self, params):
+        """Call the SABnzbd API"""
+        if self._session.closed:
+            raise SabnzbdApiException('Session already closed')
 
-    async def call(self, params):
-        api_params = {**self.default_params, **params}
+        p = {**self._default_params, **params}
         try:
-            async with aiohttp.ClientSession() as session:
-                resp = await session.get(self.api_url, params=api_params)
-                data = await resp.json()
-                if data.get('status', True) is False:
-                    self.handle_error(data, api_params)
-                else:
-                    return data
+            async with timeout(self._timeout, loop=self._session.loop):
+                async with self._session.get(self._api_url, params=p) as resp:
+                    data = await resp.json()
+                    if data.get('status', True) is False:
+                        self._handle_error(data, params)
+                    else:
+                        return data
         except aiohttp.ClientError:
             raise SabnzbdApiException('Unable to communicate with Sabnzbd API')
+        except asyncio.TimeoutError:
+            raise SabnzbdApiException('SABnzbd API request timed out')
 
     async def refresh_data(self):
+        """Refresh the cached SABnzbd queue data"""
         queue = await self.get_queue()
         history = await self.get_history()
         totals = {}
         for k in history:
             if k[-4:] == 'size':
-                totals[k] = self.convert_size(history.get(k))
+                totals[k] = self._convert_size(history.get(k))
         self.queue = {**totals, **queue}
 
     async def get_history(self):
+        """Fetch the SABnzbd history stats"""
         params = {'mode': 'history', 'limit': 1}
-        history = await self.call(params)
+        history = await self._call(params)
         return history.get('history')
 
     async def get_queue(self):
+        """Fetch the SABnzbd queue stats"""
         params = {'mode': 'queue', 'start': '0', 'limit': '10'}
-        queue = await self.call(params)
+        queue = await self._call(params)
         return queue.get('queue')
 
     async def pause_queue(self):
+        """Pause the SABnzbd client"""
         params = {'mode': 'pause'}
-        await self.call(params)
+        await self._call(params)
 
     async def resume_queue(self):
+        """Resume the SABnzbd client"""
         params = {'mode': 'resume'}
-        await self.call(params)
+        await self._call(params)
 
     async def set_speed_limit(self, speed=100):
+        """Set the download speed limit of the SABnzbd client"""
         params = {'mode': 'config', 'name': 'speedlimit', 'value': speed}
-        await self.call(params)
+        await self._call(params)
 
     async def check_available(self):
+        """Test the connection to the SABnzbd client"""
         params = {'mode': 'queue'}
-        await self.call(params)
+        await self._call(params)
         return True
 
-    def convert_size(self, size_str):
+    def _convert_size(self, size_str):
+        """Convert units to GB"""
         suffix = size_str[-1]
         if suffix == 'K':
             multiplier = 1.0 / (1024.0 * 1024.0)
@@ -78,13 +111,15 @@ def convert_size(self, size_str):
         except ValueError:
             return 0.0
 
-    def handle_error(self, data, params):
+    def _handle_error(self, data, params):
+        """Handle an error response from the SABnzbd API"""
         error = data.get('error', 'API call failed')
         mode = params.get('mode')
         raise SabnzbdApiException(error, mode=mode)
 
 
 class SabnzbdApiException(Exception):
+    """Base exception class for all SABnzbd API errors"""
     def __init__(self, message, mode=None):
         self.message = message
         self.mode = mode

setup.py

@@ -3,15 +3,16 @@
 setup(
     name='pysabnzbd',
     packages=['pysabnzbd'],
-    version='1.0.1',
+    version='1.1.0',
     description='Python wrapper for SABnzbd API',
     author='Jerad Meisner',
     author_email='jerad.meisner@gmail.com',
     url='https://github.com/jeradM/pysabnzbd',
-    download_url='https://github.com/jeradM/pysabnzbd/archive/1.0.1.tar.gz',
+    download_url='https://github.com/jeradM/pysabnzbd/archive/1.1.0.tar.gz',
     keywords=['SABnzbd'],
     install_requires=[
-        'aiohttp'
+        'aiohttp',
+        'async-timeout'
     ],
     classifiers=[]
 )