Types in poller

.vscode/cspell.json

@@ -215,6 +215,7 @@
     "ints",
     "iohttp",
     "IOHTTP",
+    "IOLRO",
     "inprogress",
     "ipconfiguration",
     "ipconfigurations",

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

@@ -37,7 +37,7 @@
 import uuid
 from typing import IO, cast, Union, Optional, AnyStr, Dict, MutableMapping
 import urllib.parse
-from typing_extensions import Protocol
+from typing_extensions import Protocol, runtime_checkable
 
 from azure.core import __version__ as azcore_version
 from azure.core.exceptions import DecodeError, raise_with_traceback
@@ -49,15 +49,17 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+@runtime_checkable
 class HTTPRequestType(Protocol):
     """Protocol compatible with new rest request and legacy transport request"""
 
     headers: MutableMapping[str, str]
     url: str
     method: str
-    body: bytes
+    body: Optional[Union[bytes, Dict[str, Union[str, int]]]]
 
 
+@runtime_checkable
 class HTTPResponseType(Protocol):
     """Protocol compatible with new rest response and legacy transport response"""
 
@@ -73,9 +75,16 @@ def status_code(self) -> int:
     def content_type(self) -> Optional[str]:
         ...
 
+    @property
+    def request(self) -> HTTPRequestType:
+        ...
+
     def text(self, encoding: Optional[str] = None) -> str:
         ...
 
+    def body(self) -> bytes:
+        ...
+
 
 class HeadersPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseType]):
     """A simple policy that sends the given headers with the request.

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

@@ -25,21 +25,24 @@
 # --------------------------------------------------------------------------
 import datetime
 import email.utils
+from typing import Optional
+
 from ...utils._utils import _FixedOffset, case_insensitive_dict
 
 
-def _parse_http_date(text):
+def _parse_http_date(text: str) -> datetime.datetime:
     """Parse a HTTP date format into datetime."""
     parsed_date = email.utils.parsedate_tz(text)
-    return datetime.datetime(*parsed_date[:6], tzinfo=_FixedOffset(parsed_date[9] / 60))
+    return datetime.datetime(*parsed_date[:6], tzinfo=_FixedOffset(parsed_date[9] / 60))  # type: ignore
 
 
-def parse_retry_after(retry_after):
+def parse_retry_after(retry_after: str) -> float:
     """Helper to parse Retry-After and get value in seconds.
 
     :param str retry_after: Retry-After header
-    :rtype: int
+    :rtype: float
     """
+    delay: float  # Using the Mypy recommendation to use float for "int or float"
     try:
         delay = int(retry_after)
     except ValueError:
@@ -49,7 +52,7 @@ def parse_retry_after(retry_after):
     return max(0, delay)
 
 
-def get_retry_after(response):
+def get_retry_after(response) -> Optional[float]:
     """Get the value of Retry-After in seconds.
 
     :param response: The PipelineResponse object

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

@@ -149,7 +149,7 @@ def open(self):
     def close(self):
         """Close the session if it is not externally owned."""
 
-    def sleep(self, duration):  # pylint: disable=no-self-use
+    def sleep(self, duration: float) -> None:  # pylint: disable=no-self-use
         time.sleep(duration)
 
 

sdk/core/azure-core/azure/core/polling/_async_poller.py

@@ -71,7 +71,12 @@ async def run(self):  # pylint:disable=invalid-overridden-method
         """
 
 
-async def async_poller(client, initial_response, deserialization_callback, polling_method):
+async def async_poller(
+    client: Any,
+    initial_response: Any,
+    deserialization_callback: Callable[[Any], PollingReturnType],
+    polling_method: AsyncPollingMethod[PollingReturnType],
+):
     """Async Poller for long running operations.
 
     .. deprecated:: 1.5.0
@@ -109,7 +114,7 @@ def __init__(
         self,
         client: Any,
         initial_response: Any,
-        deserialization_callback: Callable,
+        deserialization_callback: Callable[[Any], PollingReturnType],
         polling_method: AsyncPollingMethod[PollingReturnType],
     ):
         self._polling_method = polling_method

sdk/core/azure-core/azure/core/polling/_poller.py

@@ -64,14 +64,18 @@ def from_continuation_token(cls, continuation_token: str, **kwargs) -> Tuple[Any
         raise TypeError("Polling method '{}' doesn't support from_continuation_token".format(cls.__name__))
 
 
-class NoPolling(PollingMethod):
+class NoPolling(PollingMethod[PollingReturnType]):
     """An empty poller that returns the deserialized initial response."""
 
+    _deserialization_callback: Callable[[Any], PollingReturnType]
+    """Deserialization callback passed during initialization"""
+
     def __init__(self):
         self._initial_response = None
-        self._deserialization_callback = None
 
-    def initialize(self, _: Any, initial_response: Any, deserialization_callback: Callable) -> None:
+    def initialize(
+        self, _: Any, initial_response: Any, deserialization_callback: Callable[[Any], PollingReturnType]
+    ) -> None:
         self._initial_response = initial_response
         self._deserialization_callback = deserialization_callback
 
@@ -92,7 +96,7 @@ def finished(self) -> bool:
         """
         return True
 
-    def resource(self) -> Any:
+    def resource(self) -> PollingReturnType:
         return self._deserialization_callback(self._initial_response)
 
     def get_continuation_token(self) -> str:
@@ -130,7 +134,7 @@ def __init__(
         self,
         client: Any,
         initial_response: Any,
-        deserialization_callback: Callable,
+        deserialization_callback: Callable[[Any], PollingReturnType],
         polling_method: PollingMethod[PollingReturnType],
     ) -> None:
         self._callbacks: List[Callable] = []
@@ -147,10 +151,11 @@ def __init__(
 
         # Prepare thread execution
         self._thread = None
-        self._done = None
+        self._done = threading.Event()
         self._exception = None
-        if not self._polling_method.finished():
-            self._done = threading.Event()
+        if self._polling_method.finished():
+            self._done.set()
+        else:
             self._thread = threading.Thread(
                 target=with_current_context(self._start),
                 name="LROPoller({})".format(uuid.uuid4()),
@@ -266,7 +271,7 @@ def add_done_callback(self, func: Callable) -> None:
          argument, a completed LongRunningOperation.
         """
         # Still use "_done" and not "done", since CBs are executed inside the thread.
-        if self._done is None or self._done.is_set():
+        if self._done.is_set():
             func(self._polling_method)
         # Let's add them still, for consistency (if you wish to access to it for some reasons)
         self._callbacks.append(func)

sdk/core/azure-core/azure/core/polling/async_base_polling.py

@@ -23,24 +23,60 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
+from typing import TYPE_CHECKING, TypeVar, cast
 from ..exceptions import HttpResponseError
 from .base_polling import (
     _failed,
     BadStatus,
     BadResponse,
     OperationFailed,
-    LROBasePolling,
+    _SansIOLROBasePolling,
     _raise_if_bad_http_status_and_method,
 )
+from ._async_poller import AsyncPollingMethod
 from ..pipeline._tools import is_rest
 
+
+if TYPE_CHECKING:
+    from azure.core import AsyncPipelineClient
+    from azure.core.pipeline import PipelineResponse
+    from azure.core.pipeline.transport import AsyncHttpTransport
+    from azure.core._pipeline_client_async import _AsyncContextManagerCloseable
+    from azure.core.pipeline.policies._universal import HTTPRequestType, HTTPResponseType
+
+    AsyncHTTPResponseType = TypeVar("AsyncHTTPResponseType", bound="_AsyncContextManagerCloseable")
+    AsyncPipelineResponseType = PipelineResponse[HTTPRequestType, AsyncHTTPResponseType]
+
+PollingReturnType = TypeVar("PollingReturnType")
+
 __all__ = ["AsyncLROBasePolling"]
 
 
-class AsyncLROBasePolling(LROBasePolling):
-    """A subclass or LROBasePolling that redefine "run" as async."""
+class AsyncLROBasePolling(
+    _SansIOLROBasePolling[PollingReturnType, "AsyncPipelineClient[HTTPRequestType, AsyncHTTPResponseType]"],
+    AsyncPollingMethod[PollingReturnType],
+):
+    """A base LRO async poller.
+
+    This assumes a basic flow:
+    - I analyze the response to decide the polling approach
+    - I poll
+    - I ask the final resource depending of the polling approach
 
-    async def run(self):  # pylint:disable=invalid-overridden-method
+    If your polling need are more specific, you could implement a PollingMethod directly
+    """
+
+    _initial_response: "AsyncPipelineResponseType"
+    """Store the initial response."""
+
+    _pipeline_response: "AsyncPipelineResponseType"
+    """Store the latest received HTTP response, initialized by the first answer."""
+
+    @property
+    def _transport(self) -> "AsyncHttpTransport":
+        return self._client._pipeline._transport  # pylint: disable=protected-access
+
+    async def run(self) -> None:
         try:
             await self._poll()
 
@@ -59,7 +95,7 @@ async def run(self):  # pylint:disable=invalid-overridden-method
         except OperationFailed as err:
             raise HttpResponseError(response=self._pipeline_response.http_response, error=err)
 
-    async def _poll(self):  # pylint:disable=invalid-overridden-method
+    async def _poll(self) -> None:
         """Poll status of operation so long as operation is incomplete and
         we have an endpoint to query.
 
@@ -83,23 +119,23 @@ async def _poll(self):  # pylint:disable=invalid-overridden-method
             self._pipeline_response = await self.request_status(final_get_url)
             _raise_if_bad_http_status_and_method(self._pipeline_response.http_response)
 
-    async def _sleep(self, delay):  # pylint:disable=invalid-overridden-method
+    async def _sleep(self, delay: float) -> None:
         await self._transport.sleep(delay)
 
-    async def _delay(self):  # pylint:disable=invalid-overridden-method
+    async def _delay(self):
         """Check for a 'retry-after' header to set timeout,
         otherwise use configured timeout.
         """
         delay = self._extract_delay()
         await self._sleep(delay)
 
-    async def update_status(self):  # pylint:disable=invalid-overridden-method
+    async def update_status(self):
         """Update the current status of the LRO."""
         self._pipeline_response = await self.request_status(self._operation.get_polling_url())
         _raise_if_bad_http_status_and_method(self._pipeline_response.http_response)
         self._status = self._operation.get_status(self._pipeline_response)
 
-    async def request_status(self, status_link):  # pylint:disable=invalid-overridden-method
+    async def request_status(self, status_link: str) -> "AsyncPipelineResponseType":
         """Do a simple GET to this status link.
 
         This method re-inject 'x-ms-client-request-id'.
@@ -117,12 +153,17 @@ async def request_status(self, status_link):  # pylint:disable=invalid-overridde
             from azure.core.rest import HttpRequest as RestHttpRequest
 
             request = RestHttpRequest("GET", status_link)
-            return await self._client.send_request(request, _return_pipeline_response=True, **self._operation_config)
+            # Need a cast, as "_return_pipeline_response" mutate the return type, and that return type is not
+            # declared in the typing of "send_request"
+            return cast(
+                "AsyncPipelineResponseType",
+                await self._client.send_request(request, _return_pipeline_response=True, **self._operation_config),
+            )
         # if I am a azure.core.pipeline.transport.HttpResponse
-        request = self._client.get(status_link)
+        legacy_request = self._client.get(status_link)
 
         return await self._client._pipeline.run(  # pylint: disable=protected-access
-            request, stream=False, **self._operation_config
+            legacy_request, stream=False, **self._operation_config
         )