Add asyncio support

adamjmcgrath · adamjmcgrath · commit 57c3157a1932 · 2022-04-01T09:33:00.000+01:00
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -30,7 +30,7 @@ jobs:
       - python/install-packages:
           pkg-manager: pip-dist
           path-args: ".[test]"
-      - run: coverage run -m unittest discover
+      - run: coverage run -m unittest discover -s auth0/v3/test -t .
       - run: bash <(curl -s https://codecov.io/bash)
 
 workflows:
diff --git a/.flake8 b/.flake8
@@ -1,3 +1,3 @@
 [flake8]
-ignore = E501
+ignore = E501 F401
 max-line-length = 88
diff --git a/auth0/v3/management/asyncify.py b/auth0/v3/management/asyncify.py
@@ -0,0 +1,64 @@
+import asyncio
+
+import aiohttp
+
+from .rest_async import AsyncRestClient
+
+
+def _gen_async(client, method):
+    m = getattr(client, method)
+
+    async def closure(*args, **kwargs):
+        return await m(*args, **kwargs)
+
+    return closure
+
+
+def asyncify(cls):
+    methods = [
+        func
+        for func in dir(cls)
+        if callable(getattr(cls, func)) and not func.startswith("_")
+    ]
+
+    class AsyncClient(cls):
+        def __init__(
+            self,
+            domain,
+            token,
+            telemetry=True,
+            timeout=5.0,
+            protocol="https",
+            rest_options=None,
+        ):
+            super(AsyncClient, self).__init__(
+                domain,
+                token,
+                telemetry,
+                timeout,
+                protocol,
+                rest_options,
+            )
+            self.client = AsyncRestClient(
+                jwt=token, telemetry=telemetry, timeout=timeout, options=rest_options
+            )
+
+    class Wrapper(cls):
+        def __init__(self, *args, **kwargs):
+            super(Wrapper, self).__init__(*args, **kwargs)
+            self._async_client = AsyncClient(*args, **kwargs)
+            for method in methods:
+                setattr(self, f"{method}_async", _gen_async(self._async_client, method))
+
+        async def __aenter__(self):
+            """Automatically create and set session within context manager."""
+            async_rest_client = self._async_client.client
+            self._session = aiohttp.ClientSession()
+            async_rest_client.set_session(self._session)
+            return self
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            """Automatically close session within context manager."""
+            await self._session.close()
+
+    return Wrapper
diff --git a/auth0/v3/management/rest.py b/auth0/v3/management/rest.py
@@ -96,7 +96,7 @@ def __init__(self, jwt, telemetry=True, timeout=5.0, options=None):
             self.base_headers.update(
                 {
                     "User-Agent": "Python/{}".format(py_version),
-                    "Auth0-Client": base64.b64encode(auth0_client),
+                    "Auth0-Client": base64.b64encode(auth0_client).decode(),
                 }
             )
 
diff --git a/auth0/v3/management/rest_async.py b/auth0/v3/management/rest_async.py
@@ -0,0 +1,99 @@
+import asyncio
+import json
+
+import aiohttp
+
+from .rest import EmptyResponse, JsonResponse, PlainResponse
+from .rest import Response as _Response
+from .rest import RestClient
+
+
+def _clean_params(params):
+    if params is None:
+        return params
+    return {k: v for k, v in params.items() if v is not None}
+
+
+class AsyncRestClient(RestClient):
+    """Provides simple methods for handling all RESTful api endpoints.
+
+    Args:
+        telemetry (bool, optional): Enable or disable Telemetry
+            (defaults to True)
+        timeout (float or tuple, optional): Change the requests
+            connect and read timeout. Pass a tuple to specify
+            both values separately or a float to set both to it.
+            (defaults to 5.0 for both)
+        options (RestClientOptions): Pass an instance of
+            RestClientOptions to configure additional RestClient
+            options, such as rate-limit retries. Overrides matching
+            options passed to the constructor.
+            (defaults to 3)
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(AsyncRestClient, self).__init__(*args, **kwargs)
+        self._session = None
+
+    def set_session(self, session):
+        """Set Client Session to improve performance by reusing session.
+        Session should be closed manually or within context manager.
+        """
+        self._session = session
+
+    async def _request(self, *args, **kwargs):
+        headers = kwargs.get("headers", self.base_headers)
+        if self._session is not None:
+            # Request with re-usable session
+            async with self._session.request(
+                *args, **kwargs, headers=headers
+            ) as response:
+                return await self._process_response(response)
+        else:
+            # Request without re-usable session
+            async with aiohttp.ClientSession(headers=headers) as session:
+                async with session.request(*args, **kwargs) as response:
+                    return await self._process_response(response)
+
+    async def get(self, url, params=None):
+        return await self._request("get", url, params=_clean_params(params))
+
+    async def post(self, url, data=None):
+        return await self._request("post", url, json=data)
+
+    async def file_post(self, url, data=None, files=None):
+        headers = self.base_headers.copy()
+        headers.pop("Content-Type", None)
+        return await self._request("post", url, data={**data, **files}, headers=headers)
+
+    async def patch(self, url, data=None):
+        return await self._request("patch", url, json=data)
+
+    async def put(self, url, data=None):
+        return await self._request("put", url, json=data)
+
+    async def delete(self, url, params=None, data=None):
+        return await self._request(
+            "delete", url, json=data, params=_clean_params(params) or {}
+        )
+
+    async def _process_response(self, response):
+        parsed_response = await self._parse(response)
+        return parsed_response.content()
+
+    async def _parse(self, response):
+        text = await response.text()
+        requests_response = RequestsResponse(response, text)
+        if not text:
+            return EmptyResponse(response.status)
+        try:
+            return JsonResponse(requests_response)
+        except ValueError:
+            return PlainResponse(requests_response)
+
+
+class RequestsResponse(object):
+    def __init__(self, response, text):
+        self.status_code = response.status
+        self.headers = response.headers
+        self.text = text
diff --git a/auth0/v3/management/utils.py b/auth0/v3/management/utils.py
@@ -0,0 +1,16 @@
+import sys
+
+
+def is_async_available():
+    if sys.version_info >= (3, 0):
+        try:
+            import asyncio
+
+            import aiohttp
+            import async_timeout
+
+            return True
+        except ImportError:
+            pass
+
+    return False
diff --git a/auth0/v3/test_async/__init__.py b/auth0/v3/test_async/__init__.py
diff --git a/auth0/v3/test_async/test_asyncify.py b/auth0/v3/test_async/test_asyncify.py
@@ -0,0 +1,38 @@
+import asyncio
+import unittest
+
+import mock
+from aioresponses import aioresponses
+
+from auth0.v3.management.asyncify import asyncify
+from auth0.v3.management.clients import Clients
+
+payload = {"foo": "bar"}
+
+
+class TestAsyncify(unittest.IsolatedAsyncioTestCase):
+    @aioresponses()
+    async def test_all(self, mocked):
+        mocked.get(
+            "https://example.com/api/v2/clients?include_fields=true",
+            status=200,
+            payload=payload,
+        )
+        c = asyncify(Clients)(
+            domain="example.com",
+            token="jwt",
+        )
+        self.assertEqual(await c.all_async(), payload)
+
+    @aioresponses()
+    async def test_all_shared_session(self, mocked):
+        mocked.get(
+            "https://example.com/api/v2/clients?include_fields=true",
+            status=200,
+            payload=payload,
+        )
+        async with asyncify(Clients)(
+            domain="example.com",
+            token="jwt",
+        ) as c:
+            self.assertEqual(await c.all_async(), payload)
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,11 @@
 -e .
+aiohttp==3.8.1
+aioresponses==0.7.3
+aiosignal==1.2.0
+alabaster==0.7.12
+async-timeout==4.0.2
+attrs==21.4.0
+Authlib==1.0.0
 Babel==2.9.1
 black==22.3.0
 certifi==2021.10.8
@@ -8,30 +15,43 @@ charset-normalizer==2.0.12
 click==8.0.4
 coverage==6.3.2
 cryptography==36.0.2
+Deprecated==1.2.13
 distlib==0.3.4
 docutils==0.17.1
 filelock==3.6.0
 flake8==4.0.1
+Flask==2.0.3
+Flask-Cors==3.0.10
+frozenlist==1.3.0
 identify==2.4.12
 idna==3.3
 imagesize==1.3.0
+iniconfig==1.1.1
 isort==5.10.1
+itsdangerous==2.1.1
 Jinja2==3.1.1
+jwcrypto==1.0
 MarkupSafe==2.1.1
 mccabe==0.6.1
 mock==4.0.3
+multidict==6.0.2
 mypy-extensions==0.4.3
 nodeenv==1.6.0
 packaging==21.3
 pathspec==0.9.0
 platformdirs==2.5.1
+pluggy==1.0.0
 pre-commit==2.17.0
+py==1.11.0
 pycodestyle==2.8.0
 pycparser==2.21
 pyflakes==2.4.0
 Pygments==2.11.2
 PyJWT==2.3.0
 pyparsing==3.0.7
+pytest==7.1.0
+pytest-mock==3.7.0
+python-dotenv==0.19.2
 pytz==2022.1
 pyupgrade==2.31.1
 PyYAML==6.0
@@ -51,3 +71,6 @@ toml==0.10.2
 tomli==2.0.1
 urllib3==1.26.9
 virtualenv==20.13.4
+Werkzeug==2.0.3
+wrapt==1.14.0
+yarl==1.7.2

Original file line number	Diff line number	Diff line change
`@@ -96,7 +96,7 @@ def __init__(self, jwt, telemetry=True, timeout=5.0, options=None):`
`96`	`96`	`self.base_headers.update(`
`97`	`97`	`{`
`98`	`98`	`"User-Agent": "Python/{}".format(py_version),`
`99`		`- "Auth0-Client": base64.b64encode(auth0_client),`
	`99`	`+ "Auth0-Client": base64.b64encode(auth0_client).decode(),`
`100`	`100`	`}`
`101`	`101`	`)`
`102`	`102`