typing improvement for policies (Azure#31018)

* typing improvement for policies

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

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

Co-authored-by: Laurent Mazuel <[email protected]>

* address review feedback

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

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

* update

* update

* update

* update

* Update __init__.py

* update

* update

* update

---------

Co-authored-by: Laurent Mazuel <[email protected]>

Author: xiangyan99

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

@@ -4,8 +4,10 @@
 # license information.
 # -------------------------------------------------------------------------
 import time
-from typing import TYPE_CHECKING, Dict, Optional, TypeVar
-
+from typing import TYPE_CHECKING, Optional, TypeVar, MutableMapping
+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
 from . import HTTPPolicy, SansIOHTTPPolicy
 from ...exceptions import ServiceRequestError
 
@@ -17,10 +19,9 @@
         AzureKeyCredential,
         AzureSasCredential,
     )
-    from azure.core.pipeline import PipelineRequest, PipelineResponse
-    from azure.core.pipeline.policies._universal import HTTPRequestType
 
-HTTPResponseTypeVar = TypeVar("HTTPResponseTypeVar")
+HTTPResponseType = TypeVar("HTTPResponseType", HttpResponse, LegacyHttpResponse)
+HTTPRequestType = TypeVar("HTTPRequestType", HttpRequest, LegacyHttpRequest)
 
 
 # pylint:disable=too-few-public-methods
@@ -39,7 +40,7 @@ def __init__(self, credential: "TokenCredential", *scopes: str, **kwargs) -> Non
         self._token: Optional["AccessToken"] = None
 
     @staticmethod
-    def _enforce_https(request: "PipelineRequest") -> None:
+    def _enforce_https(request: PipelineRequest[HTTPRequestType]) -> None:
         # move 'enforce_https' from options to context so it persists
         # across retries but isn't passed to a transport implementation
         option = request.context.options.pop("enforce_https", None)
@@ -55,10 +56,10 @@ def _enforce_https(request: "PipelineRequest") -> None:
             )
 
     @staticmethod
-    def _update_headers(headers: Dict[str, str], token: str) -> None:
+    def _update_headers(headers: MutableMapping[str, str], token: str) -> None:
         """Updates the Authorization header with the bearer token.
 
-        :param dict headers: The HTTP Request headers
+        :param MutableMapping[str, str] headers: The HTTP Request headers
         :param str token: The OAuth token.
         """
         headers["Authorization"] = "Bearer {}".format(token)
@@ -68,7 +69,7 @@ def _need_new_token(self) -> bool:
         return not self._token or self._token.expires_on - time.time() < 300
 
 
-class BearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, HTTPPolicy):
+class BearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, HTTPPolicy[HTTPRequestType, HTTPResponseType]):
     """Adds a bearer token Authorization header to requests.
 
     :param credential: The credential.
@@ -77,7 +78,7 @@ class BearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, HTTPPolicy):
     :raises: :class:`~azure.core.exceptions.ServiceRequestError`
     """
 
-    def on_request(self, request: "PipelineRequest") -> None:
+    def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
         """Called before the policy sends a request.
 
         The base implementation authorizes the request with a bearer token.
@@ -90,7 +91,7 @@ def on_request(self, request: "PipelineRequest") -> 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", *scopes: str, **kwargs) -> None:
+    def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs) -> 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
@@ -102,7 +103,7 @@ def authorize_request(self, request: "PipelineRequest", *scopes: str, **kwargs)
         self._token = self._credential.get_token(*scopes, **kwargs)
         self._update_headers(request.http_request.headers, self._token.token)
 
-    def send(self, request: "PipelineRequest") -> "PipelineResponse":
+    def send(self, request: PipelineRequest[HTTPRequestType]) -> PipelineResponse[HTTPRequestType, HTTPResponseType]:
         """Authorize request with a bearer token and send it to the next policy
 
         :param request: The pipeline request object
@@ -136,7 +137,9 @@ def send(self, request: "PipelineRequest") -> "PipelineResponse":
 
         return response
 
-    def on_challenge(self, request: "PipelineRequest", response: "PipelineResponse") -> bool:
+    def on_challenge(
+        self, request: PipelineRequest[HTTPRequestType], response: PipelineResponse[HTTPRequestType, HTTPResponseType]
+    ) -> bool:
         """Authorize request according to an authentication challenge
 
         This method is called when the resource provider responds 401 with a WWW-Authenticate header.
@@ -149,7 +152,9 @@ def on_challenge(self, request: "PipelineRequest", response: "PipelineResponse")
         # pylint:disable=unused-argument
         return False
 
-    def on_response(self, request: "PipelineRequest", response: "PipelineResponse") -> None:
+    def on_response(
+        self, request: PipelineRequest[HTTPRequestType], response: PipelineResponse[HTTPRequestType, HTTPResponseType]
+    ) -> None:
         """Executed after the request comes back from the next policy.
 
         :param request: Request to be modified after returning from the policy.
@@ -158,7 +163,7 @@ def on_response(self, request: "PipelineRequest", response: "PipelineResponse")
         :type response: ~azure.core.pipeline.PipelineResponse
         """
 
-    def on_exception(self, request: "PipelineRequest") -> None:
+    def on_exception(self, request: PipelineRequest[HTTPRequestType]) -> None:
         """Executed when an exception is raised while executing the next policy.
 
         This method is executed inside the exception handler.
@@ -170,7 +175,7 @@ def on_exception(self, request: "PipelineRequest") -> None:
         return
 
 
-class AzureKeyCredentialPolicy(SansIOHTTPPolicy["HTTPRequestType", HTTPResponseTypeVar]):
+class AzureKeyCredentialPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseType]):
     """Adds a key header for the provided credential.
 
     :param credential: The credential used to authenticate requests.
@@ -199,11 +204,11 @@ def __init__(
         self._name = name
         self._prefix = prefix + " " if prefix else ""
 
-    def on_request(self, request: "PipelineRequest[HTTPRequestType]") -> None:
+    def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
         request.http_request.headers[self._name] = f"{self._prefix}{self._credential.key}"
 
 
-class AzureSasCredentialPolicy(SansIOHTTPPolicy):
+class AzureSasCredentialPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseType]):
     """Adds a shared access signature to query for the provided credential.
 
     :param credential: The credential used to authenticate requests.
@@ -217,7 +222,7 @@ def __init__(self, credential: "AzureSasCredential", **kwargs) -> None:  # pylin
             raise ValueError("credential can not be None")
         self._credential = credential
 
-    def on_request(self, request):
+    def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
         url = request.http_request.url
         query = request.http_request.query
         signature = self._credential.signature

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

@@ -5,22 +5,27 @@
 # -------------------------------------------------------------------------
 import asyncio
 import time
-from typing import TYPE_CHECKING, Any, Awaitable, Optional, cast
+from typing import TYPE_CHECKING, Any, Awaitable, Optional, cast, TypeVar
 
 from azure.core.credentials import AccessToken
+from azure.core.pipeline import PipelineRequest, PipelineResponse
 from azure.core.pipeline.policies import AsyncHTTPPolicy
 from azure.core.pipeline.policies._authentication import (
     _BearerTokenCredentialPolicyBase,
 )
+from azure.core.pipeline.transport import AsyncHttpResponse as LegacyAsyncHttpResponse, HttpRequest as LegacyHttpRequest
+from azure.core.rest import AsyncHttpResponse, HttpRequest
 
 from .._tools_async import await_result
 
 if TYPE_CHECKING:
     from azure.core.credentials_async import AsyncTokenCredential
-    from azure.core.pipeline import PipelineRequest, PipelineResponse
 
+AsyncHTTPResponseType = TypeVar("AsyncHTTPResponseType", AsyncHttpResponse, LegacyAsyncHttpResponse)
+HTTPRequestType = TypeVar("HTTPRequestType", HttpRequest, LegacyHttpRequest)
 
-class AsyncBearerTokenCredentialPolicy(AsyncHTTPPolicy):
+
+class AsyncBearerTokenCredentialPolicy(AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]):
     """Adds a bearer token Authorization header to requests.
 
     :param credential: The credential.
@@ -36,7 +41,7 @@ def __init__(self, credential: "AsyncTokenCredential", *scopes: str, **kwargs: A
         self._scopes = scopes
         self._token: Optional["AccessToken"] = None
 
-    async def on_request(self, request: "PipelineRequest") -> None:
+    async def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
         """Adds a bearer token Authorization header to request and sends request to next policy.
 
         :param request: The pipeline request object to be modified.
@@ -52,7 +57,7 @@ async def on_request(self, request: "PipelineRequest") -> None:
                     self._token = await await_result(self._credential.get_token, *self._scopes)
         request.http_request.headers["Authorization"] = "Bearer " + cast(AccessToken, self._token).token
 
-    async def authorize_request(self, request: "PipelineRequest", *scopes: str, **kwargs: Any) -> None:
+    async 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
@@ -65,7 +70,9 @@ async def authorize_request(self, request: "PipelineRequest", *scopes: str, **kw
             self._token = await await_result(self._credential.get_token, *scopes, **kwargs)
         request.http_request.headers["Authorization"] = "Bearer " + cast(AccessToken, self._token).token
 
-    async def send(self, request: "PipelineRequest") -> "PipelineResponse":
+    async def send(
+        self, request: PipelineRequest[HTTPRequestType]
+    ) -> PipelineResponse[HTTPRequestType, AsyncHTTPResponseType]:
         """Authorize request with a bearer token and send it to the next policy
 
         :param request: The pipeline request object
@@ -101,7 +108,11 @@ async def send(self, request: "PipelineRequest") -> "PipelineResponse":
 
         return response
 
-    async def on_challenge(self, request: "PipelineRequest", response: "PipelineResponse") -> bool:
+    async def on_challenge(
+        self,
+        request: PipelineRequest[HTTPRequestType],
+        response: PipelineResponse[HTTPRequestType, AsyncHTTPResponseType],
+    ) -> bool:
         """Authorize request according to an authentication challenge
 
         This method is called when the resource provider responds 401 with a WWW-Authenticate header.
@@ -114,7 +125,11 @@ async def on_challenge(self, request: "PipelineRequest", response: "PipelineResp
         # pylint:disable=unused-argument
         return False
 
-    def on_response(self, request: "PipelineRequest", response: "PipelineResponse") -> Optional[Awaitable[None]]:
+    def on_response(
+        self,
+        request: PipelineRequest[HTTPRequestType],
+        response: PipelineResponse[HTTPRequestType, AsyncHTTPResponseType],
+    ) -> Optional[Awaitable[None]]:
         """Executed after the request comes back from the next policy.
 
         :param request: Request to be modified after returning from the policy.
@@ -123,7 +138,7 @@ def on_response(self, request: "PipelineRequest", response: "PipelineResponse")
         :type response: ~azure.core.pipeline.PipelineResponse
         """
 
-    def on_exception(self, request: "PipelineRequest") -> None:
+    def on_exception(self, request: PipelineRequest[HTTPRequestType]) -> None:
         """Executed when an exception is raised while executing the next policy.
 
         This method is executed inside the exception handler.

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

@@ -29,7 +29,6 @@
 import logging
 
 from typing import (
-    TYPE_CHECKING,
     Generic,
     TypeVar,
     Union,
@@ -41,29 +40,23 @@
 
 from azure.core.pipeline import PipelineRequest, PipelineResponse
 
-if TYPE_CHECKING:
-    from azure.core.pipeline.transport import HttpTransport
-
-
-HTTPResponseTypeVar = TypeVar("HTTPResponseTypeVar")
-HTTPRequestTypeVar = TypeVar("HTTPRequestTypeVar")
+HTTPResponseType = TypeVar("HTTPResponseType")
+HTTPRequestType = TypeVar("HTTPRequestType")
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class HTTPPolicy(abc.ABC, Generic[HTTPRequestTypeVar, HTTPResponseTypeVar]):
+class HTTPPolicy(abc.ABC, Generic[HTTPRequestType, HTTPResponseType]):
     """An HTTP policy ABC.
 
     Use with a synchronous pipeline.
     """
 
-    next: "HTTPPolicy[HTTPRequestTypeVar, HTTPResponseTypeVar]"
+    next: "HTTPPolicy[HTTPRequestType, HTTPResponseType]"
     """Pointer to the next policy or a transport (wrapped as a policy). Will be set at pipeline creation."""
 
     @abc.abstractmethod
-    def send(
-        self, request: PipelineRequest[HTTPRequestTypeVar]
-    ) -> PipelineResponse[HTTPRequestTypeVar, HTTPResponseTypeVar]:
+    def send(self, request: PipelineRequest[HTTPRequestType]) -> PipelineResponse[HTTPRequestType, HTTPResponseType]:
         """Abstract send method for a synchronous pipeline. Mutates the request.
 
         Context content is dependent on the HttpTransport.
@@ -75,7 +68,7 @@ def send(
         """
 
 
-class SansIOHTTPPolicy(Generic[HTTPRequestTypeVar, HTTPResponseTypeVar]):
+class SansIOHTTPPolicy(Generic[HTTPRequestType, HTTPResponseType]):
     """Represents a sans I/O policy.
 
     SansIOHTTPPolicy is a base class for policies that only modify or
@@ -87,7 +80,7 @@ class SansIOHTTPPolicy(Generic[HTTPRequestTypeVar, HTTPResponseTypeVar]):
     but they will then be tied to AsyncPipeline usage.
     """
 
-    def on_request(self, request: PipelineRequest[HTTPRequestTypeVar]) -> Union[None, Awaitable[None]]:
+    def on_request(self, request: PipelineRequest[HTTPRequestType]) -> Union[None, Awaitable[None]]:
         """Is executed before sending the request from next policy.
 
         :param request: Request to be modified before sent from next policy.
@@ -96,8 +89,8 @@ def on_request(self, request: PipelineRequest[HTTPRequestTypeVar]) -> Union[None
 
     def on_response(
         self,
-        request: PipelineRequest[HTTPRequestTypeVar],
-        response: PipelineResponse[HTTPRequestTypeVar, HTTPResponseTypeVar],
+        request: PipelineRequest[HTTPRequestType],
+        response: PipelineResponse[HTTPRequestType, HTTPResponseType],
     ) -> Union[None, Awaitable[None]]:
         """Is executed after the request comes back from the policy.
 
@@ -109,7 +102,7 @@ def on_response(
 
     def on_exception(
         self,
-        request: PipelineRequest[HTTPRequestTypeVar],  # pylint: disable=unused-argument
+        request: PipelineRequest[HTTPRequestType],  # pylint: disable=unused-argument
     ) -> None:
         """Is executed if an exception is raised while executing the next policy.
 
@@ -129,7 +122,7 @@ def on_exception(
         return
 
 
-class RequestHistory(Generic[HTTPRequestTypeVar, HTTPResponseTypeVar]):
+class RequestHistory(Generic[HTTPRequestType, HTTPResponseType]):
     """A container for an attempted request and the applicable response.
 
     This is used to document requests/responses that resulted in redirected/retried requests.
@@ -144,12 +137,12 @@ class RequestHistory(Generic[HTTPRequestTypeVar, HTTPResponseTypeVar]):
 
     def __init__(
         self,
-        http_request: HTTPRequestTypeVar,
-        http_response: Optional[HTTPResponseTypeVar] = None,
+        http_request: HTTPRequestType,
+        http_response: Optional[HTTPResponseType] = None,
         error: Optional[Exception] = None,
         context: Optional[Dict[str, Any]] = None,
     ) -> None:
-        self.http_request = copy.deepcopy(http_request)
-        self.http_response = http_response
-        self.error = error
-        self.context = context
+        self.http_request: HTTPRequestType = copy.deepcopy(http_request)
+        self.http_response: Optional[HTTPResponseType] = http_response
+        self.error: Optional[Exception] = error
+        self.context: Optional[Dict[str, Any]] = context

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

@@ -25,32 +25,27 @@
 # --------------------------------------------------------------------------
 import abc
 
-from typing import TYPE_CHECKING, Generic, TypeVar
-
+from typing import Generic, TypeVar
 from .. import PipelineRequest, PipelineResponse
 
-if TYPE_CHECKING:
-    from ..transport._base_async import AsyncHttpTransport
-
-
-AsyncHTTPResponseTypeVar = TypeVar("AsyncHTTPResponseTypeVar")
-HTTPResponseTypeVar = TypeVar("HTTPResponseTypeVar")
-HTTPRequestTypeVar = TypeVar("HTTPRequestTypeVar")
+AsyncHTTPResponseType = TypeVar("AsyncHTTPResponseType")
+HTTPResponseType = TypeVar("HTTPResponseType")
+HTTPRequestType = TypeVar("HTTPRequestType")
 
 
-class AsyncHTTPPolicy(abc.ABC, Generic[HTTPRequestTypeVar, AsyncHTTPResponseTypeVar]):
+class AsyncHTTPPolicy(abc.ABC, Generic[HTTPRequestType, AsyncHTTPResponseType]):
     """An async HTTP policy ABC.
 
     Use with an asynchronous pipeline.
     """
 
-    next: "AsyncHTTPPolicy[HTTPRequestTypeVar, AsyncHTTPResponseTypeVar]"
+    next: "AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]"
     """Pointer to the next policy or a transport (wrapped as a policy). Will be set at pipeline creation."""
 
     @abc.abstractmethod
     async def send(
-        self, request: PipelineRequest[HTTPRequestTypeVar]
-    ) -> PipelineResponse[HTTPRequestTypeVar, AsyncHTTPResponseTypeVar]:
+        self, request: PipelineRequest[HTTPRequestType]
+    ) -> PipelineResponse[HTTPRequestType, AsyncHTTPResponseType]:
         """Abstract send method for a asynchronous pipeline. Mutates the request.
 
         Context content is dependent on the HttpTransport.