Allow session to be passed in, add async context manager

Author: jjlawren

README.md

@@ -18,10 +18,10 @@ PAYLOAD = {
 }
 
 async def main():
-    plexauth = PlexAuth(PAYLOAD)
-    await plexauth.initiate_auth()
-    print("Complete auth at URL: {}".format(plexauth.auth_url()))
-    token = await plexauth.token()
+    async with PlexAuth(PAYLOAD) as plexauth:
+        await plexauth.initiate_auth()
+        print("Complete auth at URL: {}".format(plexauth.auth_url()))
+        token = await plexauth.token()
     
     if token:
         print("Token: {}".format(token))

plexauth.py

@@ -5,29 +5,34 @@
 import urllib.parse
 import uuid
 
-__version__ = '0.0.1'
+__version__ = '0.0.3'
 
 CODES_URL = 'https://plex.tv/api/v2/pins.json?strong=true'
 AUTH_URL = 'https://app.plex.tv/auth#!?{}'
 TOKEN_URL = 'https://plex.tv/api/v2/pins/{}'
 
 class PlexAuth():
 
-    def __init__(self, payload):
+    def __init__(self, payload, session=None):
         '''Create PlexAuth instance.'''
         self._client_identifier = str(uuid.uuid4())
         self._code = None
         self._identifier = None
         self._payload = payload
         self._payload['X-Plex-Client-Identifier'] = self._client_identifier
 
+        self._local_session = False
+        self._session = session
+        if session is None:
+            self._session = aiohttp.ClientSession()
+            self._local_session = True
+
     async def initiate_auth(self):
         '''Request codes needed to create an auth URL. Starts external timeout.'''
-        async with aiohttp.ClientSession() as session:
-            async with session.post(CODES_URL, data=self._payload) as resp:
-                response = await resp.json()
-                self._code = response['code']
-                self._identifier = response['id']
+        async with self._session.post(CODES_URL, data=self._payload) as resp:
+            response = await resp.json()
+            self._code = response['code']
+            self._identifier = response['id']
 
     def auth_url(self, forward_url=None):
         '''Return an auth URL for the user to follow.'''
@@ -44,13 +49,12 @@ def auth_url(self, forward_url=None):
     async def request_auth_token(self):
         '''Request an auth token from Plex.'''
         token_url = TOKEN_URL.format(self._code)
-        async with aiohttp.ClientSession() as session:
-            payload = dict(self._payload)
-            payload['Accept'] = 'application/json'
-            async with session.get(TOKEN_URL.format(self._identifier), headers=payload) as resp:
-                response = await resp.json()
-                token = response['authToken']
-                return token
+        payload = dict(self._payload)
+        payload['Accept'] = 'application/json'
+        async with self._session.get(TOKEN_URL.format(self._identifier), headers=payload) as resp:
+            response = await resp.json()
+            token = response['authToken']
+            return token
 
     async def token(self, timeout=60):
         '''Poll Plex endpoint until a token is retrieved or times out.'''
@@ -64,3 +68,16 @@ async def token(self, timeout=60):
                 break_loop = True
 
         return token
+
+    async def close(self):
+        """Close open client session."""
+        if self._local_session:
+            await self._session.close()
+
+    async def __aenter__(self):
+        """Async enter."""
+        return self
+
+    async def __aexit__(self, *exc_info):
+        """Async exit."""
+        await self.close()

setup.py

@@ -3,7 +3,7 @@
 with open('README.md') as f:
     long_description = f.read()
 
-VERSION="0.0.1"
+VERSION="0.0.3"
 
 setup(
     name='plexauth',