improve type completeness (#33720)

* improve type completeness

* update

* update

* updates

* update

* updates

* updates

* update

* update

* enable verifytypes

Author: xiangyan99

sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py

@@ -50,11 +50,11 @@ def __init__(
         self._redirect_uri = redirect_uri
         super(AuthorizationCodeCredential, self).__init__()
 
-    def __enter__(self):
+    def __enter__(self) -> "AuthorizationCodeCredential":
         self._client.__enter__()
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: Any) -> None:
         self._client.__exit__(*args)
 
     def close(self) -> None:

sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py

@@ -60,10 +60,10 @@ def __init__(
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
         self._process_timeout = process_timeout
 
-    def __enter__(self):
+    def __enter__(self) -> "AzureCliCredential":
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: Any) -> None:
         pass
 
     def close(self) -> None:

sdk/identity/azure-identity/azure/identity/_credentials/azure_powershell.py

@@ -74,10 +74,10 @@ def __init__(
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
         self._process_timeout = process_timeout
 
-    def __enter__(self):
+    def __enter__(self) -> "AzurePowerShellCredential":
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: Any) -> None:
         pass
 
     def close(self) -> None:

sdk/identity/azure-identity/azure/identity/_credentials/chained.py

@@ -56,12 +56,12 @@ def __init__(self, *credentials):
         self._successful_credential = None  # type: Optional[TokenCredential]
         self.credentials = credentials
 
-    def __enter__(self):
+    def __enter__(self) -> "ChainedTokenCredential":
         for credential in self.credentials:
             credential.__enter__()  # type: ignore
         return self
 
-    def __exit__(self, *args: Any):
+    def __exit__(self, *args: Any) -> None:
         for credential in self.credentials:
             credential.__exit__(*args)  # type: ignore
 

sdk/identity/azure-identity/azure/identity/_credentials/client_assertion.py

@@ -55,11 +55,11 @@ def __init__(self, tenant_id: str, client_id: str, func: Callable[[], str], **kw
         )
         super(ClientAssertionCredential, self).__init__(**kwargs)
 
-    def __enter__(self):
+    def __enter__(self) -> "ClientAssertionCredential":
         self._client.__enter__()
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: Any) -> None:
         self._client.__exit__(*args)
 
     def close(self) -> None:

sdk/identity/azure-identity/azure/identity/_credentials/device_code.py

@@ -39,6 +39,7 @@ class DeviceCodeCredential(InteractiveCredential):
         - ``verification_uri`` (str) the URL the user must visit
         - ``user_code`` (str) the code the user must enter there
         - ``expires_on`` (datetime.datetime) the UTC time at which the code will expire
+
         If this argument isn't provided, the credential will print instructions to stdout.
     :paramtype prompt_callback: Callable[str, str, ~datetime.datetime]
     :keyword AuthenticationRecord authentication_record: :class:`AuthenticationRecord` returned by :func:`authenticate`

sdk/identity/azure-identity/azure/identity/_credentials/environment.py

@@ -108,12 +108,12 @@ def __init__(self, **kwargs: Any) -> None:
             else:
                 _LOGGER.info("No environment configuration found.")
 
-    def __enter__(self):
+    def __enter__(self) -> "EnvironmentCredential":
         if self._credential:
             self._credential.__enter__()
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: Any) -> None:
         if self._credential:
             self._credential.__exit__(*args)
 

sdk/identity/azure-identity/azure/identity/_credentials/imds.py

@@ -65,7 +65,7 @@ def __init__(self, **kwargs: Any) -> None:
             self._endpoint_available = None
         self._user_assigned_identity = "client_id" in kwargs or "identity_config" in kwargs
 
-    def __enter__(self):
+    def __enter__(self) -> "ImdsCredential":
         self._client.__enter__()
         return self
 

sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py

@@ -96,12 +96,14 @@ def __init__(self, **kwargs: Any) -> None:
             _LOGGER.info("%s will use IMDS", self.__class__.__name__)
             self._credential = ImdsCredential(**kwargs)
 
-    def __enter__(self):
-        self._credential.__enter__()
+    def __enter__(self) -> "ManagedIdentityCredential":
+        if self._credential:
+            self._credential.__enter__()  # type: ignore
         return self
 
-    def __exit__(self, *args):
-        self._credential.__exit__(*args)
+    def __exit__(self, *args: Any) -> None:
+        if self._credential:
+            self._credential.__exit__(*args)  # type: ignore
 
     def close(self) -> None:
         """Close the credential's transport session."""

sdk/identity/azure-identity/azure/identity/_credentials/shared_cache.py

@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, TypeVar
 from azure.core.credentials import AccessToken
 
 from .silent import SilentAuthenticationCredential
@@ -16,6 +16,9 @@
     from azure.core.credentials import TokenCredential
 
 
+T = TypeVar("T", bound="_SharedTokenCacheCredential")
+
+
 class SharedTokenCacheCredential:
     """Authenticates using tokens in the local cache shared between Microsoft applications.
 
@@ -40,12 +43,12 @@ def __init__(self, username: Optional[str] = None, **kwargs: Any) -> None:
         else:
             self._credential = _SharedTokenCacheCredential(username=username, **kwargs)
 
-    def __enter__(self):
-        self._credential.__enter__()
+    def __enter__(self) -> "SharedTokenCacheCredential":
+        self._credential.__enter__()  # type: ignore
         return self
 
-    def __exit__(self, *args):
-        self._credential.__exit__(*args)
+    def __exit__(self, *args: Any) -> None:
+        self._credential.__exit__(*args)  # type: ignore
 
     def close(self) -> None:
         """Close the credential's transport session."""
@@ -97,9 +100,9 @@ def supported() -> bool:
 class _SharedTokenCacheCredential(SharedTokenCacheBase):
     """The original SharedTokenCacheCredential, which doesn't use msal.ClientApplication"""
 
-    def __enter__(self):
+    def __enter__(self: T) -> T:
         if self._client:
-            self._client.__enter__()
+            self._client.__enter__()  # type: ignore
         return self
 
     def __exit__(self, *args):