Type Fix for identity folder (#33479)

* fix aml_on_behalf_of

* fix misc

* fix misc - 2

* fix managed_identity_client

* fix misc - 3

* update pyproject.toml - 1

* update pyproject.toml - 2

Author: CoderKevinZhang

sdk/ml/azure-ai-ml/azure/__init__.py

@@ -9,4 +9,4 @@
 * Autorest objects.
 * SDK v2 support the new YAML file format.
 """
-__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

sdk/ml/azure-ai-ml/azure/ai/__init__.py

@@ -2,4 +2,4 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

sdk/ml/azure-ai-ml/azure/ai/ml/__init__.py

@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)
 
 import logging
 

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_aio/_credentials/aml_on_behalf_of.py

@@ -4,7 +4,9 @@
 
 import functools
 import os
+from typing import Any, Optional
 
+from azure.core.credentials import AccessToken
 from azure.core.pipeline.transport import HttpRequest
 
 from .._internal import AsyncContextManager
@@ -22,10 +24,10 @@ class AzureMLOnBehalfOfCredential(AsyncContextManager):
     """
     # pylint: enable=line-too-long
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any):
         self._credential = _AzureMLOnBehalfOfCredential(**kwargs)
 
-    async def get_token(self, *scopes, **kwargs):
+    async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
         """Request an access token for `scopes`.
 
         This method is called automatically by Azure SDK clients.
@@ -38,12 +40,12 @@ async def get_token(self, *scopes, **kwargs):
 
         return await self._credential.get_token(*scopes, **kwargs)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "AzureMLOnBehalfOfCredential":
         if self._credential:
             await self._credential.__aenter__()
         return self
 
-    async def close(self):
+    async def close(self) -> None:
         """Close the credential's transport session."""
         if self._credential:
             await self._credential.__aexit__()
@@ -77,7 +79,7 @@ def _get_client_args(**kwargs):
 
 
 def _get_request(url, resource):
-    # type: (str, str, dict) -> HttpRequest
+    # type: (str, str) -> HttpRequest
     request = HttpRequest("GET", url)
     request.format_parameters(dict({"resource": resource}))
     return request

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_aio/_internal/__init__.py

@@ -2,17 +2,18 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 import abc
+from typing import Any
 
 
 class AsyncContextManager(abc.ABC):
     @abc.abstractmethod
-    async def close(self):
+    async def close(self) -> None:
         pass
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "AsyncContextManager":
         return self
 
-    async def __aexit__(self, *args):
+    async def __aexit__(self, *args: Any) -> None:
         await self.close()
 
 

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_aio/_internal/get_token_mixin.py

@@ -25,7 +25,7 @@ def __init__(self, *args: "Any", **kwargs: "Any") -> None:
         self._last_request_time = 0
 
         # https://github.com/python/mypy/issues/5887
-        super(GetTokenMixin, self).__init__(*args, **kwargs)  # type: ignore
+        super(GetTokenMixin, self).__init__(*args, **kwargs)
 
     @abc.abstractmethod
     # pylint: disable-next=docstring-missing-param

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_aio/_internal/managed_identity_base.py

@@ -3,40 +3,38 @@
 # ---------------------------------------------------------
 
 import abc
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from ..._exceptions import CredentialUnavailableError
 from . import AsyncContextManager
 from .get_token_mixin import GetTokenMixin
 from .managed_identity_client import AsyncManagedIdentityClient
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
-
     from azure.core.credentials import AccessToken
 
 
 class AsyncManagedIdentityBase(AsyncContextManager, GetTokenMixin):
     """Base class for internal credentials using AsyncManagedIdentityClient."""
 
-    def __init__(self, **kwargs: "Any") -> None:
+    def __init__(self, **kwargs: Any) -> None:
         super().__init__()
         self._client = self.get_client(**kwargs)
 
     @abc.abstractmethod
-    def get_client(self, **kwargs: "Any") -> "Optional[AsyncManagedIdentityClient]":
+    def get_client(self, **kwargs: Any) -> "Optional[AsyncManagedIdentityClient]":
         pass
 
     @abc.abstractmethod
     def get_unavailable_message(self) -> str:
         pass
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "AsyncManagedIdentityBase":
         if self._client:
             await self._client.__aenter__()
         return self
 
-    async def __aexit__(self, *args):
+    async def __aexit__(self, *args: Any) -> None:
         if self._client:
             await self._client.__aexit__(*args)
 

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_aio/_internal/managed_identity_client.py

@@ -20,7 +20,7 @@
 
 # pylint:disable=async-client-bad-name,missing-client-constructor-parameter-credential
 class AsyncManagedIdentityClient(AsyncContextManager, ManagedIdentityClientBase):
-    async def __aenter__(self):
+    async def __aenter__(self) -> "AsyncManagedIdentityClient":
         await self._pipeline.__aenter__()
         return self
 
@@ -30,7 +30,8 @@ async def close(self) -> None:
     async def request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         # pylint:disable=invalid-overridden-method
         resource = _scopes_to_resource(*scopes)
-        request = self._request_factory(resource, self._identity_config)  # pylint: disable=no-member
+        # pylint: disable=no-member
+        request = self._request_factory(resource, self._identity_config)  # type: ignore
         request_time = int(time.time())
         response = await self._pipeline.run(request, retry_on_methods=[request.method], **kwargs)
         token = self._process_response(response=response, request_time=request_time, resource=resource)

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_credentials/_AzureMLSparkOnBehalfOfCredential.py

@@ -5,7 +5,7 @@
 
 import functools
 import os
-from typing import Optional
+from typing import Any, Optional
 
 from azure.core.pipeline.transport import HttpRequest
 
@@ -14,7 +14,7 @@
 
 
 class _AzureMLSparkOnBehalfOfCredential(ManagedIdentityBase):
-    def get_client(self, **kwargs) -> Optional[ManagedIdentityClient]:
+    def get_client(self, **kwargs: Any) -> Optional[ManagedIdentityClient]:
         client_args = _get_client_args(**kwargs)
         if client_args:
             return ManagedIdentityClient(**client_args)
@@ -24,7 +24,7 @@ def get_unavailable_message(self) -> str:
         return "AzureML Spark On Behalf of credentials not available in this environment"
 
 
-def _get_client_args(**kwargs) -> Optional[dict]:
+def _get_client_args(**kwargs: Any) -> Optional[dict]:
     # Override default settings if provided via arguments
     if len(kwargs) > 0:
         env_key_from_kwargs = [
@@ -80,7 +80,7 @@ def _get_client_args(**kwargs) -> Optional[dict]:
     )
 
 
-def _get_request(url, resource) -> HttpRequest:
+def _get_request(url: str, resource: Any) -> HttpRequest:
     obo_access_token = os.environ.get("AZUREML_OBO_CANARY_TOKEN")
     experiment_name = os.environ.get("AZUREML_ARM_PROJECT_NAME")
     run_id = os.environ.get("AZUREML_RUN_ID")

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_credentials/aml_on_behalf_of.py

@@ -4,12 +4,14 @@
 
 import functools
 import os
+from typing import Any, Optional, Union
 
+from azure.core.credentials import AccessToken
 from azure.core.pipeline.transport import HttpRequest
 
-from ._AzureMLSparkOnBehalfOfCredential import _AzureMLSparkOnBehalfOfCredential
 from .._internal.managed_identity_base import ManagedIdentityBase
 from .._internal.managed_identity_client import ManagedIdentityClient
+from ._AzureMLSparkOnBehalfOfCredential import _AzureMLSparkOnBehalfOfCredential
 
 
 class AzureMLOnBehalfOfCredential(object):
@@ -23,14 +25,16 @@ class AzureMLOnBehalfOfCredential(object):
     """
     # pylint: enable=line-too-long
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any):
         provider_type = os.environ.get("AZUREML_DATAPREP_TOKEN_PROVIDER")
+        self._credential: Union[_AzureMLSparkOnBehalfOfCredential, _AzureMLOnBehalfOfCredential]
+
         if provider_type == "sparkobo":  # cspell:disable-line
             self._credential = _AzureMLSparkOnBehalfOfCredential(**kwargs)
         else:
             self._credential = _AzureMLOnBehalfOfCredential(**kwargs)
 
-    def get_token(self, *scopes, **kwargs):
+    def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
         """Request an access token for `scopes`.
 
         This method is called automatically by Azure SDK clients.
@@ -43,11 +47,11 @@ def get_token(self, *scopes, **kwargs):
 
         return self._credential.get_token(*scopes, **kwargs)
 
-    def __enter__(self):
+    def __enter__(self) -> "AzureMLOnBehalfOfCredential":
         self._credential.__enter__()
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: Any) -> None:
         self._credential.__exit__(*args)
 
     def close(self):
@@ -84,7 +88,7 @@ def _get_client_args(**kwargs):
 
 
 def _get_request(url, resource):
-    # type: (str, str, dict) -> HttpRequest
+    # type: (str, str) -> HttpRequest
     request = HttpRequest("GET", url)
     request.format_parameters(dict({"resource": resource}))
     return request