Core typings (#28114)

* Core typings * Restore mypy clean * Black it * Mypy+PyLint+Black * Remove object inheritance * dict/list are not typing before 3.9 * Runtime fixes * Forgotten List * Update _authentication.py * Replace more dict by Dict * More typing fixes * More Dict fixes * Fix List typing * Pyright Paging * Simplify some ABC import * Pylint clean-up * Fix typing for Identity * Keep 'next' in paging as this used in some Py3 code... * Form can have int values as well * Better annotations * Pylint has trouble on 3.7 * Type configuration * MatchCondition doc * PyLint happyness * Feedback from Kashif * Redesign some of the async pipeline types * Fix for 3.7 * Remove unecessary type * Feedback * Male kwarg only syntax * Correct enum doc * Full typing for CloudEvent * New feedback from Johan * More feedback * Feedback * Enabling mypy and typecheck * Improve policy typing * Improve response typing * Update sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py Co-authored-by: Anna Tisch <[email protected]> * Update sdk/core/azure-core/azure/core/_pipeline_client_async.py Co-authored-by: Anna Tisch <[email protected]> * Update sdk/core/azure-core/azure/core/_pipeline_client_async.py Co-authored-by: Anna Tisch <[email protected]> * Feedback * Black * Enum typing * Remove init for better typing * Fix typing * PyLint --------- Co-authored-by: Anna Tisch <[email protected]>
Azure · Jan 30, 2023 · 8176b0f · 8176b0f
1 parent e00bfa2
commit 8176b0f
Show file tree

Hide file tree

Showing 43 changed files with 913 additions and 873 deletions.
diff --git a/sdk/core/azure-core/azure/core/_enum_meta.py b/sdk/core/azure-core/azure/core/_enum_meta.py
@@ -23,8 +23,8 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-
-from enum import EnumMeta
+from typing import Any
+from enum import EnumMeta, Enum
 
 
 class CaseInsensitiveEnumMeta(EnumMeta):
@@ -43,13 +43,13 @@ class MyCustomEnum(str, Enum, metaclass=CaseInsensitiveEnumMeta):
 
     """
 
-    def __getitem__(cls, name):
+    def __getitem__(cls, name: str) -> Any:
         # disabling pylint bc of pylint bug https://github.com/PyCQA/astroid/issues/713
         return super(  # pylint: disable=no-value-for-parameter
             CaseInsensitiveEnumMeta, cls
         ).__getitem__(name.upper())
 
-    def __getattr__(cls, name):
+    def __getattr__(cls, name: str) -> Enum:
         """Return the enum member matching `name`
         We use __getattr__ instead of descriptors or inserting into the enum
         class' __dict__ in order to support `name` and `value` being both

diff --git a/sdk/core/azure-core/azure/core/_match_conditions.py b/sdk/core/azure-core/azure/core/_match_conditions.py
@@ -30,8 +30,17 @@
 class MatchConditions(Enum):
     """An enum to describe match conditions."""
 
-    Unconditionally = 1  # Matches any condition
-    IfNotModified = 2  # If the target object is not modified. Usually it maps to etag=<specific etag>
-    IfModified = 3  # Only if the target object is modified. Usually it maps to etag!=<specific etag>
-    IfPresent = 4  # If the target object exists. Usually it maps to etag='*'
-    IfMissing = 5  # If the target object does not exist. Usually it maps to etag!='*'
+    Unconditionally = 1
+    """Matches any condition"""
+
+    IfNotModified = 2
+    """If the target object is not modified. Usually it maps to etag=<specific etag>"""
+
+    IfModified = 3
+    """Only if the target object is modified. Usually it maps to etag!=<specific etag>"""
+
+    IfPresent = 4
+    """If the target object exists. Usually it maps to etag='*'"""
+
+    IfMissing = 5
+    """If the target object does not exist. Usually it maps to etag!='*'"""
diff --git a/sdk/core/azure-core/azure/core/_pipeline_client.py b/sdk/core/azure-core/azure/core/_pipeline_client.py
@@ -27,7 +27,6 @@
 import logging
 from collections.abc import Iterable
 from typing import (
-    Any,
     TypeVar,
     TYPE_CHECKING,
 )
@@ -169,8 +168,7 @@ def _build_pipeline(self, config, **kwargs):  # pylint: disable=no-self-use
 
         return Pipeline(transport, policies)
 
-    def send_request(self, request, **kwargs):
-        # type: (HTTPRequestType, Any) -> HTTPResponseType
+    def send_request(self, request: "HTTPRequestType", **kwargs) -> "HTTPResponseType":
         """Method that runs the network request through the client's chained policies.
 
         >>> from azure.core.rest import HttpRequest

diff --git a/sdk/core/azure-core/azure/core/_pipeline_client_async.py b/sdk/core/azure-core/azure/core/_pipeline_client_async.py
@@ -26,7 +26,16 @@
 
 import logging
 import collections.abc
-from typing import Any, Awaitable, TypeVar
+from typing import (
+    Any,
+    Awaitable,
+    TypeVar,
+    AsyncContextManager,
+    Generator,
+    cast,
+    TYPE_CHECKING,
+)
+from typing_extensions import Protocol
 from .configuration import Configuration
 from .pipeline import AsyncPipeline
 from .pipeline.transport._base import PipelineClientBase
@@ -38,33 +47,88 @@
     AsyncRetryPolicy,
 )
 
+
+if TYPE_CHECKING:  # Protocol and non-Protocol can't mix in Python 3.7
+
+    class _AsyncContextManagerCloseable(AsyncContextManager, Protocol):
+        """Defines a context manager that is closeable at the same time."""
+
+        async def close(self):
+            ...
+
+
 HTTPRequestType = TypeVar("HTTPRequestType")
-AsyncHTTPResponseType = TypeVar("AsyncHTTPResponseType")
+AsyncHTTPResponseType = TypeVar(
+    "AsyncHTTPResponseType", bound="_AsyncContextManagerCloseable"
+)
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class _AsyncContextManager(collections.abc.Awaitable):
-    def __init__(self, wrapped: collections.abc.Awaitable):
+class _Coroutine(Awaitable[AsyncHTTPResponseType]):
+    """Wrapper to get both context manager and awaitable in place.
+
+    Naming it "_Coroutine" because if you don't await it makes the error message easier:
+    >>> result = client.send_request(request)
+    >>> result.text()
+    AttributeError: '_Coroutine' object has no attribute 'text'
+
+    Indeed, the message for calling a coroutine without waiting would be:
+    AttributeError: 'coroutine' object has no attribute 'text'
+
+    This allows the dev to either use the "async with" syntax, or simply the object directly.
+    It's also why "send_request" is not declared as async, since it couldn't be both easily.
+
+    "wrapped" must be an awaitable that returns an object that:
+    - has an async "close()"
+    - has an "__aexit__" method (IOW, is an async context manager)
+
+    This permits this code to work for both requests.
+
+    ```python
+    from azure.core import AsyncPipelineClient
+    from azure.core.rest import HttpRequest
+
+    async def main():
+
+        request = HttpRequest("GET", "https://httpbin.org/user-agent")
+        async with AsyncPipelineClient("https://httpbin.org/") as client:
+            # Can be used directly
+            result = await client.send_request(request)
+            print(result.text())
+
+            # Can be used as an async context manager
+            async with client.send_request(request) as result:
+                print(result.text())
+    ```
+
+    :param wrapped: Must be an awaitable the returns an async context manager that supports async "close()"
+    """
+
+    def __init__(self, wrapped: Awaitable[AsyncHTTPResponseType]) -> None:
         super().__init__()
-        self.wrapped = wrapped
-        self.response = None
+        self._wrapped = wrapped
+        # If someone tries to use the object without awaiting, they will get a
+        # AttributeError: '_Coroutine' object has no attribute 'text'
+        self._response: AsyncHTTPResponseType = cast(AsyncHTTPResponseType, None)
 
-    def __await__(self):
-        return self.wrapped.__await__()
+    def __await__(self) -> Generator[Any, None, AsyncHTTPResponseType]:
+        return self._wrapped.__await__()
 
-    async def __aenter__(self):
-        self.response = await self
-        return self.response
+    async def __aenter__(self) -> AsyncHTTPResponseType:
+        self._response = await self
+        return self._response
 
-    async def __aexit__(self, *args):
-        await self.response.__aexit__(*args)
+    async def __aexit__(self, *args) -> None:
+        await self._response.__aexit__(*args)
 
-    async def close(self):
-        await self.response.close()
+    async def close(self) -> None:
+        await self._response.close()
 
 
-class AsyncPipelineClient(PipelineClientBase):
+class AsyncPipelineClient(
+    PipelineClientBase, AsyncContextManager["AsyncPipelineClient"]
+):
     """Service client core methods.
 
     Builds an AsyncPipeline client.
@@ -212,4 +276,4 @@ def send_request(
         :rtype: ~azure.core.rest.AsyncHttpResponse
         """
         wrapped = self._make_pipeline_call(request, stream=stream, **kwargs)
-        return _AsyncContextManager(wrapped=wrapped)
+        return _Coroutine(wrapped=wrapped)
diff --git a/sdk/core/azure-core/azure/core/async_paging.py b/sdk/core/azure-core/azure/core/async_paging.py
@@ -33,6 +33,7 @@
     Tuple,
     Optional,
     Awaitable,
+    Any,
 )
 
 from .exceptions import AzureError
@@ -85,10 +86,10 @@ def __init__(
         self._extract_data = extract_data
         self.continuation_token = continuation_token
         self._did_a_call_already = False
-        self._response = None
-        self._current_page = None
+        self._response: Optional[ResponseType] = None
+        self._current_page: Optional[AsyncIterator[ReturnType]] = None
 
-    async def __anext__(self):
+    async def __anext__(self) -> AsyncIterator[ReturnType]:
         if self.continuation_token is None and self._did_a_call_already:
             raise StopAsyncIteration("End of paging")
         try:
@@ -112,18 +113,16 @@ async def __anext__(self):
 
 
 class AsyncItemPaged(AsyncIterator[ReturnType]):
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Return an async iterator of items.
 
         args and kwargs will be passed to the AsyncPageIterator constructor directly,
         except page_iterator_class
         """
         self._args = args
         self._kwargs = kwargs
-        self._page_iterator = (
-            None
-        )  # type: Optional[AsyncIterator[AsyncIterator[ReturnType]]]
-        self._page = None  # type: Optional[AsyncIterator[ReturnType]]
+        self._page_iterator: Optional[AsyncIterator[AsyncIterator[ReturnType]]] = None
+        self._page: Optional[AsyncIterator[ReturnType]] = None
         self._page_iterator_class = self._kwargs.pop(
             "page_iterator_class", AsyncPageIterator
         )

diff --git a/sdk/core/azure-core/azure/core/configuration.py b/sdk/core/azure-core/azure/core/configuration.py
@@ -23,9 +23,10 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
+from typing import Union, Optional
 
 
-class Configuration(object):
+class Configuration:
     """Provides the home for all of the configurable policies in the pipeline.
 
     A new Configuration object provides no default policies and does not specify in what
@@ -88,16 +89,17 @@ def __init__(self, **kwargs):
         self.polling_interval = kwargs.get("polling_interval", 30)
 
 
-class ConnectionConfiguration(object):
+class ConnectionConfiguration:
     """HTTP transport connection configuration settings.
 
     Common properties that can be configured on all transports. Found in the
     Configuration object.
 
-    :keyword int connection_timeout: A single float in seconds for the connection timeout. Defaults to 300 seconds.
-    :keyword int read_timeout: A single float in seconds for the read timeout. Defaults to 300 seconds.
-    :keyword bool connection_verify: SSL certificate verification. Enabled by default. Set to False to disable,
+    :keyword float connection_timeout: A single float in seconds for the connection timeout. Defaults to 300 seconds.
+    :keyword float read_timeout: A single float in seconds for the read timeout. Defaults to 300 seconds.
+    :keyword connection_verify: SSL certificate verification. Enabled by default. Set to False to disable,
      alternatively can be set to the path to a CA_BUNDLE file or directory with certificates of trusted CAs.
+    :paramtype connection_verify: bool or str
     :keyword str connection_cert: Client-side certificates. You can specify a local cert to use as client side
      certificate, as a single file (containing the private key and the certificate) or as a tuple of both files' paths.
     :keyword int connection_data_block_size: The block size of data sent over the connection. Defaults to 4096 bytes.
@@ -112,9 +114,18 @@ class ConnectionConfiguration(object):
             :caption: Configuring transport connection settings.
     """
 
-    def __init__(self, **kwargs):
-        self.timeout = kwargs.pop("connection_timeout", 300)
-        self.read_timeout = kwargs.pop("read_timeout", 300)
-        self.verify = kwargs.pop("connection_verify", True)
-        self.cert = kwargs.pop("connection_cert", None)
-        self.data_block_size = kwargs.pop("connection_data_block_size", 4096)
+    def __init__(
+        self,  # pylint: disable=unused-argument
+        *,
+        connection_timeout: float = 300,
+        read_timeout: float = 300,
+        connection_verify: Union[bool, str] = True,
+        connection_cert: Optional[str] = None,
+        connection_data_block_size: int = 4096,
+        **kwargs
+    ) -> None:
+        self.timeout = connection_timeout
+        self.read_timeout = read_timeout
+        self.verify = connection_verify
+        self.cert = connection_cert
+        self.data_block_size = connection_data_block_size