Refactor OAuth flows to use OAuth2 as a base class.

httpx_auth/_oauth2/authentication_responses_server.py

@@ -166,7 +166,7 @@ def handle_timeout(self) -> None:
         raise TimeoutOccurred(self.timeout)
 
 
-def request_new_grant(grant_details: GrantDetails) -> (str, str):
+def request_new_grant(grant_details: GrantDetails) -> tuple[str, str]:
     """
     Ask for a new OAuth2 grant.
     :return: A tuple (state, grant)

httpx_auth/_oauth2/authorization_code.py

@@ -1,5 +1,5 @@
 from hashlib import sha512
-from typing import Generator, Iterable, Union
+from typing import Iterable, Union
 
 import httpx
 
@@ -8,14 +8,14 @@
 from httpx_auth._oauth2.browser import BrowserAuth
 from httpx_auth._oauth2.common import (
     request_new_grant_with_post,
-    OAuth2,
+    OAuth2BaseAuth,
     _add_parameters,
     _pop_parameter,
     _get_query_parameter,
 )
 
 
-class OAuth2AuthorizationCode(httpx.Auth, SupportMultiAuth, BrowserAuth):
+class OAuth2AuthorizationCode(OAuth2BaseAuth, SupportMultiAuth, BrowserAuth):
     """
     Authorization Code Grant
 
@@ -70,6 +70,7 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
             raise Exception("Token URL is mandatory.")
 
         BrowserAuth.__init__(self, kwargs)
+        OAuth2BaseAuth.__init__(self)
 
         self.header_name = kwargs.pop("header_name", None) or "Authorization"
         self.header_value = kwargs.pop("header_value", None) or "Bearer {token}"
@@ -129,17 +130,8 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
         self.refresh_data = {"grant_type": "refresh_token"}
         self.refresh_data.update(kwargs)
 
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=self.request_new_token,
-            on_expired_token=self.refresh_token,
-        )
+    def _update_user_request(self, request: httpx.Request, token: str) -> None:
         request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
 
     def request_new_token(self) -> tuple:
         # Request code

httpx_auth/_oauth2/authorization_code_pkce.py

@@ -1,7 +1,6 @@
 import base64
 import os
 from hashlib import sha256, sha512
-from typing import Generator
 
 import httpx
 
@@ -10,13 +9,13 @@
 from httpx_auth._oauth2.browser import BrowserAuth
 from httpx_auth._oauth2.common import (
     request_new_grant_with_post,
-    OAuth2,
+    OAuth2BaseAuth,
     _add_parameters,
     _pop_parameter,
 )
 
 
-class OAuth2AuthorizationCodePKCE(httpx.Auth, SupportMultiAuth, BrowserAuth):
+class OAuth2AuthorizationCodePKCE(OAuth2BaseAuth, SupportMultiAuth, BrowserAuth):
     """
     Proof Key for Code Exchange
 
@@ -69,6 +68,7 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
             raise Exception("Token URL is mandatory.")
 
         BrowserAuth.__init__(self, kwargs)
+        OAuth2BaseAuth.__init__(self)
 
         self.client = kwargs.pop("client", None)
 
@@ -139,17 +139,8 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
         self.refresh_data = {"grant_type": "refresh_token"}
         self.refresh_data.update(kwargs)
 
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=self.request_new_token,
-            on_expired_token=self.refresh_token,
-        )
+    def _update_user_request(self, request: httpx.Request, token: str) -> None:
         request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
 
     def request_new_token(self) -> tuple:
         # Request code

httpx_auth/_oauth2/client_credentials.py

@@ -1,16 +1,16 @@
 from hashlib import sha512
-from typing import Generator, Union, Iterable
+from typing import Union, Iterable
 
 import httpx
 from httpx_auth._authentication import SupportMultiAuth
 from httpx_auth._oauth2.common import (
-    OAuth2,
+    OAuth2BaseAuth,
     request_new_grant_with_post,
     _add_parameters,
 )
 
 
-class OAuth2ClientCredentials(httpx.Auth, SupportMultiAuth):
+class OAuth2ClientCredentials(OAuth2BaseAuth, SupportMultiAuth):
     """
     Client Credentials Grant
 
@@ -49,6 +49,8 @@ def __init__(self, token_url: str, client_id: str, client_secret: str, **kwargs)
         if not self.client_secret:
             raise Exception("client_secret is mandatory.")
 
+        super().__init__()
+
         self.header_name = kwargs.pop("header_name", None) or "Authorization"
         self.header_value = kwargs.pop("header_value", None) or "Bearer {token}"
         if "{token}" not in self.header_value:
@@ -72,16 +74,8 @@ def __init__(self, token_url: str, client_id: str, client_secret: str, **kwargs)
         all_parameters_in_url = _add_parameters(self.token_url, self.data)
         self.state = sha512(all_parameters_in_url.encode("unicode_escape")).hexdigest()
 
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=self.request_new_token,
-        )
+    def _update_user_request(self, request: httpx.Request, token: str) -> None:
         request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
 
     def request_new_token(self) -> tuple:
         client = self.client or httpx.Client()

httpx_auth/_oauth2/common.py

@@ -1,4 +1,5 @@
-from typing import Optional
+import abc
+from typing import Callable, Generator, Optional, Union
 from urllib.parse import parse_qs, urlsplit, urlunsplit, urlencode
 
 import httpx
@@ -86,3 +87,32 @@ def request_new_grant_with_post(
 class OAuth2:
     token_cache = TokenMemoryCache()
     display = DisplaySettings()
+
+
+class OAuth2BaseAuth(abc.ABC, httpx.Auth):
+    state: Optional[str] = None
+    early_expiry: float
+
+    refresh_token: Optional[Callable]
+
+    def auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response, None]:
+        token = OAuth2.token_cache.get_token(
+            self.state,
+            early_expiry=self.early_expiry,
+            on_missing_token=self.request_new_token,
+            on_expired_token=(
+                self.refresh_token if "refresh_token" in dir(self) else None
+            ),
+        )
+        self._update_user_request(request, token)
+        yield request
+
+    @abc.abstractmethod
+    def request_new_token(self) -> Union[tuple[str, str], tuple[str, str, int]]:
+        pass  # pragma: no cover
+
+    @abc.abstractmethod
+    def _update_user_request(self, request: httpx.Request, token: str) -> None:
+        pass  # pragma: no cover

httpx_auth/_oauth2/implicit.py

@@ -1,21 +1,20 @@
 import uuid
 from hashlib import sha512
-from typing import Generator
 
 import httpx
 
 from httpx_auth._authentication import SupportMultiAuth
 from httpx_auth._oauth2 import authentication_responses_server
 from httpx_auth._oauth2.browser import BrowserAuth
 from httpx_auth._oauth2.common import (
-    OAuth2,
+    OAuth2BaseAuth,
     _add_parameters,
     _pop_parameter,
     _get_query_parameter,
 )
 
 
-class OAuth2Implicit(httpx.Auth, SupportMultiAuth, BrowserAuth):
+class OAuth2Implicit(OAuth2BaseAuth, SupportMultiAuth, BrowserAuth):
     """
     Implicit Grant
 
@@ -61,6 +60,7 @@ def __init__(self, authorization_url: str, **kwargs):
             raise Exception("Authorization URL is mandatory.")
 
         BrowserAuth.__init__(self, kwargs)
+        OAuth2BaseAuth.__init__(self)
 
         self.header_name = kwargs.pop("header_name", None) or "Authorization"
         self.header_value = kwargs.pop("header_value", None) or "Bearer {token}"
@@ -104,17 +104,11 @@ def __init__(self, authorization_url: str, **kwargs):
             self.redirect_uri_port,
         )
 
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=authentication_responses_server.request_new_grant,
-            grant_details=self.grant_details,
-        )
+    def _update_user_request(self, request: httpx.Request, token: str) -> None:
         request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
+
+    def request_new_token(self) -> tuple[str, str]:
+        return authentication_responses_server.request_new_grant(self.grant_details)
 
 
 class AzureActiveDirectoryImplicit(OAuth2Implicit):

httpx_auth/_oauth2/resource_owner_password.py

@@ -1,16 +1,15 @@
 from hashlib import sha512
-from typing import Generator
 
 import httpx
 from httpx_auth._authentication import SupportMultiAuth
 from httpx_auth._oauth2.common import (
-    OAuth2,
+    OAuth2BaseAuth,
     request_new_grant_with_post,
     _add_parameters,
 )
 
 
-class OAuth2ResourceOwnerPasswordCredentials(httpx.Auth, SupportMultiAuth):
+class OAuth2ResourceOwnerPasswordCredentials(OAuth2BaseAuth, SupportMultiAuth):
     """
     Resource Owner Password Credentials Grant
 
@@ -42,6 +41,8 @@ def __init__(self, token_url: str, username: str, password: str, **kwargs):
         Use it to provide a custom proxying rule for instance.
         :param kwargs: all additional authorization parameters that should be put as body parameters in the token URL.
         """
+        super().__init__()
+
         self.token_url = token_url
         if not self.token_url:
             raise Exception("Token URL is mandatory.")
@@ -85,17 +86,8 @@ def __init__(self, token_url: str, username: str, password: str, **kwargs):
         all_parameters_in_url = _add_parameters(self.token_url, self.data)
         self.state = sha512(all_parameters_in_url.encode("unicode_escape")).hexdigest()
 
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=self.request_new_token,
-            on_expired_token=self.refresh_token,
-        )
+    def _update_user_request(self, request: httpx.Request, token: str) -> None:
         request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
 
     def request_new_token(self) -> tuple:
         client = self.client or httpx.Client()