[req-changes] Added docstrings, fixed imports, pinned deps

Author: Aryamanz29

pyatlan/client/aio/client.py

@@ -41,6 +41,7 @@
 from pyatlan.client.aio.file import AsyncFileClient
 from pyatlan.client.aio.group import AsyncGroupClient
 from pyatlan.client.aio.impersonate import AsyncImpersonationClient
+from pyatlan.client.aio.oauth import AsyncOAuthTokenManager
 from pyatlan.client.aio.open_lineage import AsyncOpenLineageClient
 from pyatlan.client.aio.query import AsyncQueryClient
 from pyatlan.client.aio.role import AsyncRoleClient
@@ -134,13 +135,10 @@ class AsyncAtlanClient(AtlanClient):
     def __init__(self, **kwargs):
         # Initialize sync client (handles all validation, env vars, etc.)
         super().__init__(**kwargs)
-
         if self.oauth_client_id and self.oauth_client_secret and self.api_key is None:
             LOGGER.debug(
                 "API Key not provided. Using Async OAuth flow for authentication"
             )
-            from pyatlan.client.aio.oauth import AsyncOAuthTokenManager
-
             if self._oauth_token_manager:
                 LOGGER.debug("Sync oauth flow open. Closing it for Async oauth flow")
                 self._oauth_token_manager.close()

pyatlan/client/aio/oauth.py

@@ -3,6 +3,7 @@
 import asyncio
 import time
 from typing import Optional
+from urllib.parse import urljoin
 
 import httpx
 from authlib.oauth2.rfc6749 import OAuth2Token
@@ -12,6 +13,19 @@
 
 
 class AsyncOAuthTokenManager:
+    """
+    Manages OAuth tokens for asynchronous HTTP clients.
+
+    :param base_url: Base URL of the Atlan tenant.
+    :param client_id: OAuth client ID.
+    :param client_secret: OAuth client secret.
+    :param http_client: Optional asynchronous HTTP client to use.
+    :param connect_timeout: Timeout for establishing connections.
+    :param read_timeout: Timeout for reading data.
+    :param write_timeout: Timeout for writing data.
+    :param pool_timeout: Timeout for acquiring a connection from the pool.
+    """
+
     def __init__(
         self,
         base_url: str,
@@ -20,6 +34,8 @@ def __init__(
         http_client: Optional[httpx.AsyncClient] = None,
         connect_timeout: float = 30.0,
         read_timeout: float = 900.0,
+        write_timeout: float = 30.0,
+        pool_timeout: float = 30.0,
     ):
         self.base_url = base_url
         self.client_id = client_id
@@ -28,13 +44,19 @@ def __init__(
         self._lock = asyncio.Lock()
         self._http_client = http_client or httpx.AsyncClient(
             timeout=httpx.Timeout(
-                connect=connect_timeout, read=read_timeout, write=30.0, pool=30.0
+                connect=connect_timeout,
+                read=read_timeout,
+                write=write_timeout,
+                pool=pool_timeout,
             )
         )
         self._token: Optional[OAuth2Token] = None
         self._owns_client = http_client is None
 
     async def get_token(self) -> str:
+        """
+        Retrieves a valid OAuth token, refreshing it if necessary.
+        """
         async with self._lock:
             if self._token and not self._token.is_expired():
                 return str(self._token["access_token"])
@@ -73,18 +95,25 @@ async def get_token(self) -> str:
             return access_token
 
     async def invalidate_token(self):
+        """
+        Invalidates the current OAuth token.
+        """
         async with self._lock:
             self._token = None
 
     def _create_path(self, api: API):
-        from urllib.parse import urljoin
-
+        """
+        Creates the full URL for the given API endpoint.
+        """
         if self.base_url == "INTERNAL":
             return urljoin(api.endpoint.service, api.path)
         else:
             base_with_prefix = urljoin(self.base_url, api.endpoint.prefix)
             return urljoin(base_with_prefix, api.path)
 
     async def aclose(self):
+        """
+        Closes the underlying HTTP client if owned by this manager.
+        """
         if self._owns_client:
             await self._http_client.aclose()

pyatlan/client/constants.py

@@ -89,7 +89,7 @@
     WHOAMI_API, HTTPMethod.GET, HTTPStatus.OK, endpoint=EndPoint.HERACLES
 )
 
-# oauth client authentinatication
+# oauth client authentication
 GET_OAUTH_CLIENT = API(
     "oauth-clients/token",
     HTTPMethod.POST,

pyatlan/client/oauth.py

@@ -3,6 +3,7 @@
 import threading
 import time
 from typing import Optional
+from urllib.parse import urljoin
 
 import httpx
 from authlib.oauth2.rfc6749 import OAuth2Token
@@ -12,6 +13,18 @@
 
 
 class OAuthTokenManager:
+    """
+    Manages OAuth tokens for HTTP clients.
+    :param base_url: Base URL of the Atlan tenant.
+    :param client_id: OAuth client ID.
+    :param client_secret: OAuth client secret.
+    :param http_client: Optional HTTP client to use.
+    :param connect_timeout: Timeout for establishing connections.
+    :param read_timeout: Timeout for reading data.
+    :param write_timeout: Timeout for writing data.
+    :param pool_timeout: Timeout for acquiring a connection from the pool.
+    """
+
     def __init__(
         self,
         base_url: str,
@@ -20,6 +33,8 @@ def __init__(
         http_client: Optional[httpx.Client] = None,
         connect_timeout: float = 30.0,
         read_timeout: float = 900.0,
+        write_timeout: float = 30.0,
+        pool_timeout: float = 30.0,
     ):
         self.base_url = base_url
         self.client_id = client_id
@@ -28,13 +43,19 @@ def __init__(
         self._lock = threading.Lock()
         self._http_client = http_client or httpx.Client(
             timeout=httpx.Timeout(
-                connect=connect_timeout, read=read_timeout, write=30.0, pool=30.0
+                connect=connect_timeout,
+                read=read_timeout,
+                write=write_timeout,
+                pool=pool_timeout,
             )
         )
         self._token: Optional[OAuth2Token] = None
         self._owns_client = http_client is None
 
     def get_token(self) -> str:
+        """
+        Retrieves a valid OAuth token, refreshing it if necessary.
+        """
         with self._lock:
             if self._token and not self._token.is_expired():
                 return str(self._token["access_token"])
@@ -73,18 +94,25 @@ def get_token(self) -> str:
             return access_token
 
     def invalidate_token(self):
+        """
+        Invalidates the current OAuth token.
+        """
         with self._lock:
             self._token = None
 
     def _create_path(self, api: API):
-        from urllib.parse import urljoin
-
+        """
+        Creates the full URL for the given API endpoint.
+        """
         if self.base_url == "INTERNAL":
             return urljoin(api.endpoint.service, api.path)
         else:
             base_with_prefix = urljoin(self.base_url, api.endpoint.prefix)
             return urljoin(base_with_prefix, api.path)
 
     def close(self):
+        """
+        Closes the underlying HTTP client if owned by this manager.
+        """
         if self._owns_client:
             self._http_client.close()

pyatlan/pkg/utils.py

@@ -76,7 +76,7 @@ def get_client(
     This will use an API token if found in ATLAN_API_KEY, and will fallback to attempting to impersonate a user if
     ATLAN_API_KEY is empty.
 
-    :param impersonate_user_id: unique identifier (GUID) of a user or API token to impersonate
+    :param impersonate_user_id: unique identifier (GUID) of a user or API token to impersonate (default is None)
     :param set_pkg_headers: whether to set package headers on the client (default is False)
     :returns: an initialized client
     """
@@ -131,7 +131,7 @@ async def get_client_async(
     This will use an API token if found in ATLAN_API_KEY, and will fallback to attempting to impersonate a user if
     ATLAN_API_KEY is empty.
 
-    :param impersonate_user_id: unique identifier (GUID) of a user or API token to impersonate
+    :param impersonate_user_id: unique identifier (GUID) of a user or API token to impersonate (default is None)
     :param set_pkg_headers: whether to set package headers on the client (default is False)
     :returns: an initialized async client
     """

pyatlan/test_utils/__init__.py

@@ -107,7 +107,10 @@ def create_connection(
 ) -> Connection:
     admin_role_guid = str(client.role_cache.get_id_for_name("$admin"))
     to_create = Connection.create(
-        client=client, name=name, connector_type=connector_type, admin_roles=[admin_role_guid]
+        client=client,
+        name=name,
+        connector_type=connector_type,
+        admin_roles=[admin_role_guid],
     )
     response = client.asset.save(to_create)
     result = response.assets_created(asset_type=Connection)[0]

pyproject.toml

@@ -51,7 +51,7 @@ dev = [
     "mypy~=1.18.0",
     "ruff~=0.14.5",
     "types-setuptools~=80.9.0.20250822",
-    "types-Authlib",
+    "types-Authlib~=1.6.5.20251005",
     "pytest~=8.4.2",
     "pytest-vcr~=1.0.2",
     "vcrpy~=7.0.0",