Azure-Core type complete (Azure#31502)

* Pipeline complete

* Missing type

* Mypy fixes part1

* Mypy fixes part 2

* mypy fixes part 3

* mypy fixes part 4

* mypy fixes part 5

* mypy fixes part 6

* Pyright fixes

* CI fixes

* pyright fixes again

* mypy clean, no warnings

* pylint

* next-mypy

* Update sdk/core/azure-core/azure/core/pipeline/policies/_redirect.py

Co-authored-by: Krista Pratico <[email protected]>

* Update sdk/core/azure-core/azure/core/pipeline/policies/_retry.py

Co-authored-by: Krista Pratico <[email protected]>

* Move pylint ignore

* Correct aexit

* Black

* Feedback from Anna

* Kashif's feedback

* CI feedback

---------

Co-authored-by: Krista Pratico <[email protected]>

Author: lmazuel

sdk/core/azure-core/azure/core/_pipeline_client.py

@@ -23,14 +23,10 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-
+from __future__ import annotations
 import logging
 from collections.abc import Iterable
-from typing import (
-    TypeVar,
-    Generic,
-    Optional,
-)
+from typing import TypeVar, Generic, Optional, Any
 from .configuration import Configuration
 from .pipeline import Pipeline
 from .pipeline.transport._base import PipelineClientBase
@@ -82,34 +78,34 @@ def __init__(
         base_url: str,
         *,
         pipeline: Optional[Pipeline[HTTPRequestType, HTTPResponseType]] = None,
-        config: Optional[Configuration] = None,
-        **kwargs
+        config: Optional[Configuration[HTTPRequestType, HTTPResponseType]] = None,
+        **kwargs: Any,
     ):
         super(PipelineClient, self).__init__(base_url)
-        self._config: Configuration = config or Configuration(**kwargs)
+        self._config: Configuration[HTTPRequestType, HTTPResponseType] = config or Configuration(**kwargs)
         self._base_url = base_url
 
         self._pipeline = pipeline or self._build_pipeline(self._config, **kwargs)
 
-    def __enter__(self):
+    def __enter__(self) -> PipelineClient[HTTPRequestType, HTTPResponseType]:
         self._pipeline.__enter__()
         return self
 
-    def __exit__(self, *exc_details):
+    def __exit__(self, *exc_details: Any) -> None:
         self._pipeline.__exit__(*exc_details)
 
-    def close(self):
+    def close(self) -> None:
         self.__exit__()
 
     def _build_pipeline(
         self,
-        config: Configuration,
+        config: Configuration[HTTPRequestType, HTTPResponseType],
         *,
         transport: Optional[HttpTransport[HTTPRequestType, HTTPResponseType]] = None,
         policies=None,
         per_call_policies=None,
         per_retry_policies=None,
-        **kwargs
+        **kwargs,
     ) -> Pipeline[HTTPRequestType, HTTPResponseType]:
         per_call_policies = per_call_policies or []
         per_retry_policies = per_retry_policies or []
@@ -182,7 +178,7 @@ def _build_pipeline(
 
         return Pipeline(transport, policies)
 
-    def send_request(self, request: HTTPRequestType, **kwargs) -> HTTPResponseType:
+    def send_request(self, request: HTTPRequestType, **kwargs: Any) -> HTTPResponseType:
         """Method that runs the network request through the client's chained policies.
 
         >>> from azure.core.rest import HttpRequest

sdk/core/azure-core/azure/core/_pipeline_client_async.py

@@ -23,7 +23,7 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-
+from __future__ import annotations
 import logging
 import collections.abc
 from typing import (
@@ -34,8 +34,10 @@
     Generator,
     Generic,
     Optional,
+    Type,
     cast,
 )
+from types import TracebackType
 from .configuration import Configuration
 from .pipeline import AsyncPipeline
 from .pipeline.transport._base import PipelineClientBase
@@ -108,8 +110,13 @@ async def __aenter__(self) -> AsyncHTTPResponseType:
         self._response = await self
         return self._response
 
-    async def __aexit__(self, *args) -> None:
-        await self._response.__aexit__(*args)
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]] = None,
+        exc_value: Optional[BaseException] = None,
+        traceback: Optional[TracebackType] = None,
+    ) -> None:
+        await self._response.__aexit__(exc_type, exc_value, traceback)
 
 
 class AsyncPipelineClient(
@@ -150,26 +157,37 @@ def __init__(
         base_url: str,
         *,
         pipeline: Optional[AsyncPipeline[HTTPRequestType, AsyncHTTPResponseType]] = None,
-        config: Optional[Configuration] = None,
-        **kwargs
+        config: Optional[Configuration[HTTPRequestType, AsyncHTTPResponseType]] = None,
+        **kwargs: Any,
     ):
         super(AsyncPipelineClient, self).__init__(base_url)
-        self._config: Configuration = config or Configuration(**kwargs)
+        self._config: Configuration[HTTPRequestType, AsyncHTTPResponseType] = config or Configuration(**kwargs)
         self._base_url = base_url
         self._pipeline = pipeline or self._build_pipeline(self._config, **kwargs)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> AsyncPipelineClient[HTTPRequestType, AsyncHTTPResponseType]:
         await self._pipeline.__aenter__()
         return self
 
-    async def __aexit__(self, *args):
-        await self.close()
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]] = None,
+        exc_value: Optional[BaseException] = None,
+        traceback: Optional[TracebackType] = None,
+    ) -> None:
+        await self._pipeline.__aexit__(exc_type, exc_value, traceback)
 
-    async def close(self):
-        await self._pipeline.__aexit__()
+    async def close(self) -> None:
+        await self.__aexit__()
 
     def _build_pipeline(
-        self, config: Configuration, *, policies=None, per_call_policies=None, per_retry_policies=None, **kwargs
+        self,
+        config: Configuration[HTTPRequestType, AsyncHTTPResponseType],
+        *,
+        policies=None,
+        per_call_policies=None,
+        per_retry_policies=None,
+        **kwargs,
     ) -> AsyncPipeline[HTTPRequestType, AsyncHTTPResponseType]:
         transport = kwargs.get("transport")
         per_call_policies = per_call_policies or []

sdk/core/azure-core/azure/core/credentials.py

@@ -39,6 +39,7 @@ def get_token(
         :rtype: AccessToken
         :return: An AccessToken instance containing the token string and its expiration time in Unix time.
         """
+        ...
 
 
 AzureNamedKey = namedtuple("AzureNamedKey", ["name", "key"])

sdk/core/azure-core/azure/core/pipeline/__init__.py

@@ -24,11 +24,16 @@
 #
 # --------------------------------------------------------------------------
 
-from typing import TypeVar, Generic, Dict, Any, Tuple, List, Optional, overload
+from typing import TypeVar, Generic, Dict, Any, Tuple, List, Optional, overload, TYPE_CHECKING, Union
 
 HTTPResponseType = TypeVar("HTTPResponseType", covariant=True)  # pylint: disable=typevar-name-incorrect-variance
 HTTPRequestType = TypeVar("HTTPRequestType", covariant=True)  # pylint: disable=typevar-name-incorrect-variance
 
+if TYPE_CHECKING:
+    from .transport import HttpTransport, AsyncHttpTransport
+
+    TransportType = Union[HttpTransport[Any, Any], AsyncHttpTransport[Any, Any]]
+
 
 class PipelineContext(Dict[str, Any]):
     """A context object carried by the pipeline request and response containers.
@@ -45,8 +50,10 @@ class PipelineContext(Dict[str, Any]):
 
     _PICKLE_CONTEXT = {"deserialized_data"}
 
-    def __init__(self, transport: Any, **kwargs: Any) -> None:  # pylint: disable=super-init-not-called
-        self.transport: Optional[Any] = transport
+    def __init__(
+        self, transport: Optional["TransportType"], **kwargs: Any
+    ) -> None:  # pylint: disable=super-init-not-called
+        self.transport: Optional["TransportType"] = transport
         self.options = kwargs
         self._protected = ["transport", "options"]
 

sdk/core/azure-core/azure/core/pipeline/_base.py

@@ -23,10 +23,9 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-
+from __future__ import annotations
 import logging
-from typing import Generic, TypeVar, Union, Any, List, Dict, Optional, Iterable
-from contextlib import AbstractContextManager
+from typing import Generic, TypeVar, Union, Any, List, Dict, Optional, Iterable, ContextManager
 from azure.core.pipeline import (
     PipelineRequest,
     PipelineResponse,
@@ -39,9 +38,7 @@
 HTTPResponseType = TypeVar("HTTPResponseType")
 HTTPRequestType = TypeVar("HTTPRequestType")
 
-
 _LOGGER = logging.getLogger(__name__)
-PoliciesType = Iterable[Union[HTTPPolicy, SansIOHTTPPolicy]]
 
 
 def cleanup_kwargs_for_transport(kwargs: Dict[str, str]) -> None:
@@ -124,7 +121,7 @@ def send(self, request: PipelineRequest[HTTPRequestType]) -> PipelineResponse[HT
         )
 
 
-class Pipeline(AbstractContextManager, Generic[HTTPRequestType, HTTPResponseType]):
+class Pipeline(ContextManager["Pipeline"], Generic[HTTPRequestType, HTTPResponseType]):
     """A pipeline implementation.
 
     This is implemented as a context manager, that will activate the context
@@ -147,7 +144,13 @@ class Pipeline(AbstractContextManager, Generic[HTTPRequestType, HTTPResponseType
     def __init__(
         self,
         transport: HttpTransport[HTTPRequestType, HTTPResponseType],
-        policies: Optional[PoliciesType] = None,
+        policies: Optional[
+            Iterable[
+                Union[
+                    HTTPPolicy[HTTPRequestType, HTTPResponseType], SansIOHTTPPolicy[HTTPRequestType, HTTPResponseType]
+                ]
+            ]
+        ] = None,
     ) -> None:
         self._impl_policies: List[HTTPPolicy[HTTPRequestType, HTTPResponseType]] = []
         self._transport = transport
@@ -162,11 +165,11 @@ def __init__(
         if self._impl_policies:
             self._impl_policies[-1].next = _TransportRunner(self._transport)
 
-    def __enter__(self) -> "Pipeline":
+    def __enter__(self) -> Pipeline[HTTPRequestType, HTTPResponseType]:
         self._transport.__enter__()
         return self
 
-    def __exit__(self, *exc_details):  # pylint: disable=arguments-differ
+    def __exit__(self, *exc_details: Any) -> None:  # pylint: disable=arguments-differ
         self._transport.__exit__(*exc_details)
 
     @staticmethod

sdk/core/azure-core/azure/core/pipeline/_base_async.py

@@ -23,8 +23,10 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-from typing import Any, Union, Generic, TypeVar, List, Dict, Optional, Iterable
-from contextlib import AbstractAsyncContextManager
+from __future__ import annotations
+from types import TracebackType
+from typing import Any, Union, Generic, TypeVar, List, Dict, Optional, Iterable, Type
+from typing_extensions import AsyncContextManager
 
 from azure.core.pipeline import PipelineRequest, PipelineResponse, PipelineContext
 from azure.core.pipeline.policies import AsyncHTTPPolicy, SansIOHTTPPolicy
@@ -34,7 +36,6 @@
 
 AsyncHTTPResponseType = TypeVar("AsyncHTTPResponseType")
 HTTPRequestType = TypeVar("HTTPRequestType")
-AsyncPoliciesType = Iterable[Union[AsyncHTTPPolicy, SansIOHTTPPolicy]]
 
 
 class _SansIOAsyncHTTPPolicyRunner(
@@ -107,7 +108,7 @@ async def send(
         )
 
 
-class AsyncPipeline(AbstractAsyncContextManager, Generic[HTTPRequestType, AsyncHTTPResponseType]):
+class AsyncPipeline(AsyncContextManager["AsyncPipeline"], Generic[HTTPRequestType, AsyncHTTPResponseType]):
     """Async pipeline implementation.
 
     This is implemented as a context manager, that will activate the context
@@ -130,7 +131,14 @@ class AsyncPipeline(AbstractAsyncContextManager, Generic[HTTPRequestType, AsyncH
     def __init__(
         self,
         transport: AsyncHttpTransport[HTTPRequestType, AsyncHTTPResponseType],
-        policies: Optional[AsyncPoliciesType] = None,
+        policies: Optional[
+            Iterable[
+                Union[
+                    AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType],
+                    SansIOHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType],
+                ]
+            ]
+        ] = None,
     ) -> None:
         self._impl_policies: List[AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]] = []
         self._transport = transport
@@ -145,12 +153,17 @@ def __init__(
         if self._impl_policies:
             self._impl_policies[-1].next = _AsyncTransportRunner(self._transport)
 
-    async def __aenter__(self) -> "AsyncPipeline":
+    async def __aenter__(self) -> AsyncPipeline[HTTPRequestType, AsyncHTTPResponseType]:
         await self._transport.__aenter__()
         return self
 
-    async def __aexit__(self, *exc_details):  # pylint: disable=arguments-differ
-        await self._transport.__aexit__(*exc_details)
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]] = None,
+        exc_value: Optional[BaseException] = None,
+        traceback: Optional[TracebackType] = None,
+    ) -> None:
+        await self._transport.__aexit__(exc_type, exc_value, traceback)
 
     async def _prepare_multipart_mixed_request(self, request: HTTPRequestType) -> None:
         """Will execute the multipart policies.

sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py

@@ -99,7 +99,7 @@ def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
                 self._token = self._credential.get_token(*self._scopes)
         self._update_headers(request.http_request.headers, self._token.token)
 
-    def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs) -> None:
+    def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs: Any) -> None:
         """Acquire a token from the credential and authorize the request with it.
 
         Keyword arguments are passed to the credential's get_token method. The token will be cached and used to
@@ -195,13 +195,13 @@ class AzureKeyCredentialPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseTyp
     :raises: ValueError or TypeError
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=unused-argument
         self,
         credential: "AzureKeyCredential",
         name: str,
         *,
         prefix: Optional[str] = None,
-        **kwargs,  # pylint: disable=unused-argument
+        **kwargs: Any,
     ) -> None:
         super().__init__()
         if not hasattr(credential, "key"):
@@ -226,7 +226,7 @@ class AzureSasCredentialPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseTyp
     :raises: ValueError or TypeError
     """
 
-    def __init__(self, credential: "AzureSasCredential", **kwargs) -> None:  # pylint: disable=unused-argument
+    def __init__(self, credential: "AzureSasCredential", **kwargs: Any) -> None:  # pylint: disable=unused-argument
         super(AzureSasCredentialPolicy, self).__init__()
         if not credential:
             raise ValueError("credential can not be None")

sdk/core/azure-core/azure/core/pipeline/policies/_custom_hook.py

@@ -23,7 +23,7 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-from typing import TypeVar
+from typing import TypeVar, Any
 from azure.core.pipeline import PipelineRequest, PipelineResponse
 from azure.core.pipeline.transport import HttpResponse as LegacyHttpResponse, HttpRequest as LegacyHttpRequest
 from azure.core.rest import HttpResponse, HttpRequest
@@ -41,7 +41,7 @@ class CustomHookPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseType]):
     :keyword callback raw_response_hook: Callback function. Will be invoked on response.
     """
 
-    def __init__(self, **kwargs):  # pylint: disable=unused-argument,super-init-not-called
+    def __init__(self, **kwargs: Any):  # pylint: disable=unused-argument,super-init-not-called
         self._request_callback = kwargs.get("raw_request_hook")
         self._response_callback = kwargs.get("raw_response_hook")