[EH] improve typing (#33338)

* remove py2 typehints

* pyright

* adding typing

* more

* changes

* changes

* fix

* fix

* pylint

* remove

* ah

* changes

* update

* pylint

* client base changes

* typing

* adding more typing

* remove ignore pyamqp

* pull swathi changes

* fix sample errors

* pausing here

* fixes

* checkpointstore typing

* pylint

* connstate

* undo samples typing

* samples

* changes

* update

* fwd ref

* settler

* try this

* change where protocol is called from?

* revert amqp typing

* swathi changes

* remove

* emove

* initial pass

* nit

* update exceptions

* add back msg_backcompat for errors

* revert changes

* pylint

* align with swathis pyamqp

* mising default

* pyamqp error

* update error

* update

* pylint

* Update sdk/eventhub/azure-eventhub-checkpointstoreblob-aio/azure/eventhub/extensions/checkpointstoreblobaio/_blobstoragecsaio.py

Co-authored-by: Kashif Khan <[email protected]>

* pr comments

* update

* kwargs.pop

* fixing broken tests

* remove type ignore

---------

Co-authored-by: Kashif Khan <[email protected]>
Co-authored-by: Kashif Khan <[email protected]>

Author: l0lawrence

sdk/eventhub/azure-eventhub-checkpointstoreblob-aio/azure/eventhub/extensions/checkpointstoreblobaio/_blobstoragecsaio.py

@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
-from typing import Iterable, Dict, Any, Optional
+from typing import Iterable, Dict, Any, Optional, Union, TYPE_CHECKING
 import logging
 import copy
 from collections import defaultdict
@@ -13,6 +13,9 @@
 from ._vendor.storage.blob.aio import ContainerClient, BlobClient
 from ._vendor.storage.blob._shared.base_client import parse_connection_str
 
+if TYPE_CHECKING:
+    from azure.core.credentials_async import AsyncTokenCredential
+    from azure.core.credentials import AzureSasCredential, AzureNamedKeyCredential
 
 logger = logging.getLogger(__name__)
 UPLOAD_DATA = ""
@@ -41,11 +44,17 @@ class BlobCheckpointStore(CheckpointStore):
      The hostname of the secondary endpoint.
     """
 
-    def __init__(self, blob_account_url, container_name, *, credential=None, **kwargs):
-        # type(str, str, Optional[Any], Any) -> None
+    def __init__(
+        self,
+        blob_account_url: str,
+        container_name: str,
+        *,
+        credential: Optional[Union["AsyncTokenCredential", "AzureNamedKeyCredential", "AzureSasCredential"]] = None,
+        api_version: str = '2019-07-07',
+        **kwargs: Any
+    ) -> None:
         self._container_client = kwargs.pop("container_client", None)
         if not self._container_client:
-            api_version = kwargs.pop("api_version", None)
             if api_version:
                 headers = kwargs.get("headers")
                 if headers:
@@ -59,7 +68,12 @@ def __init__(self, blob_account_url, container_name, *, credential=None, **kwarg
 
     @classmethod
     def from_connection_string(
-        cls, conn_str: str, container_name: str, *, credential: Optional[Any] = None, **kwargs: Any
+        cls,
+        conn_str: str,
+        container_name: str,
+        *,
+        credential: Optional[Union["AsyncTokenCredential", "AzureNamedKeyCredential", "AzureSasCredential"]] = None,
+        **kwargs: Any
     ) -> "BlobCheckpointStore":
         """Create BlobCheckpointStore from a storage connection string.
 
@@ -88,11 +102,11 @@ def from_connection_string(
 
         return cls(account_url, container_name, credential=credential, **kwargs)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "BlobCheckpointStore":
         await self._container_client.__aenter__()
         return self
 
-    async def __aexit__(self, *args):
+    async def __aexit__(self, *args: Any) -> None:
         await self._container_client.__aexit__(*args)
 
     def _get_blob_client(self, blob_name: str) -> BlobClient:
@@ -291,8 +305,8 @@ async def update_checkpoint(self, checkpoint: Dict[str, Any], **kwargs: Any) ->
             )
 
     async def list_checkpoints(
-        self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs
-    ):
+        self, fully_qualified_namespace: str, eventhub_name: str, consumer_group: str, **kwargs: Any
+    ) -> Iterable[Dict[str, Any]]:
         """List the updated checkpoints from the storage blob.
 
         :param str fully_qualified_namespace: The fully qualified namespace that the Event Hub belongs to.

sdk/eventhub/azure-eventhub-checkpointstoreblob/azure/eventhub/extensions/checkpointstoreblob/_blobstoragecs.py

@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
-from typing import Dict, Optional, Any, Iterable, Union
+from typing import Dict, Optional, Any, Iterable, Union, TYPE_CHECKING
 import logging
 import time
 import calendar
@@ -16,6 +16,9 @@
 from ._vendor.storage.blob import BlobClient, ContainerClient
 from ._vendor.storage.blob._shared.base_client import parse_connection_str
 
+if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential, AzureSasCredential, AzureNamedKeyCredential
+
 logger = logging.getLogger(__name__)
 UPLOAD_DATA = ""
 
@@ -63,8 +66,14 @@ class BlobCheckpointStore(CheckpointStore):
 
     """
 
-    def __init__(self, blob_account_url, container_name, credential=None, **kwargs):
-        # type(str, str, Optional[Any], Any) -> None
+    def __init__(
+            self,
+            blob_account_url: str,
+            container_name: str,
+            credential: Optional[Union["AzureNamedKeyCredential", "AzureSasCredential", "TokenCredential"]] = None,
+            api_version: str = '2019-07-07',
+            **kwargs: Any
+    ) -> None:
         self._container_client = kwargs.pop("container_client", None)
         if not self._container_client:
             api_version = kwargs.pop("api_version", None)
@@ -81,9 +90,12 @@ def __init__(self, blob_account_url, container_name, credential=None, **kwargs):
 
     @classmethod
     def from_connection_string(
-        cls, conn_str, container_name, credential=None, **kwargs
-    ):
-        # type: (str, str, Optional[Any], Any) -> BlobCheckpointStore
+        cls,
+        conn_str: str,
+        container_name: str,
+        credential: Optional[Union["AzureNamedKeyCredential", "AzureSasCredential", "TokenCredential"]] = None,
+        **kwargs: Any
+    ) -> "BlobCheckpointStore":
         """Create BlobCheckpointStore from a storage connection string.
 
         :param str conn_str:
@@ -111,11 +123,11 @@ def from_connection_string(
 
         return cls(account_url, container_name, credential=credential, **kwargs)
 
-    def __enter__(self):
+    def __enter__(self) -> "BlobCheckpointStore":
         self._container_client.__enter__()
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: Any) -> None:
         self._container_client.__exit__(*args)
 
     def _get_blob_client(self, blob_name):
@@ -183,8 +195,13 @@ def _claim_one_partition(self, ownership, **kwargs):
             )
             return updated_ownership  # Keep the ownership if an unexpected error happens
 
-    def list_ownership(self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs):
-        # type: (str, str, str, Any) -> Iterable[Dict[str, Any]]
+    def list_ownership(
+        self,
+        fully_qualified_namespace: str,
+        eventhub_name: str,
+        consumer_group: str,
+        **kwargs: Any
+    ) -> Iterable[Dict[str, Any]]:
         """Retrieves a complete ownership list from the storage blob.
 
         :param str fully_qualified_namespace: The fully qualified namespace that the Event Hub belongs to.
@@ -238,8 +255,11 @@ def list_ownership(self, fully_qualified_namespace, eventhub_name, consumer_grou
             )
             raise
 
-    def claim_ownership(self, ownership_list, **kwargs):
-        # type: (Iterable[Dict[str, Any]], Any) -> Iterable[Dict[str, Any]]
+    def claim_ownership(
+        self,
+        ownership_list: Iterable[Dict[str, Any]],
+        **kwargs: Any
+    ) -> Iterable[Dict[str, Any]]:
         """Tries to claim ownership for a list of specified partitions.
 
         :param iterable[dict[str, any]] ownership_list: Iterable of dictionaries containing all the ownerships to claim.
@@ -265,8 +285,7 @@ def claim_ownership(self, ownership_list, **kwargs):
                 pass
         return gathered_results
 
-    def update_checkpoint(self, checkpoint, **kwargs):
-        # type: (Dict[str, Optional[Union[str, int]]], Any) -> None
+    def update_checkpoint(self, checkpoint: Dict[str, Union[str, int]], **kwargs: Any) -> None:
         """Updates the checkpoint using the given information for the offset, associated partition and
         consumer group in the storage blob.
 
@@ -309,9 +328,8 @@ def update_checkpoint(self, checkpoint, **kwargs):
             )
 
     def list_checkpoints(
-        self, fully_qualified_namespace, eventhub_name, consumer_group, **kwargs
-    ):
-        # type: (str, str, str, Any) -> Iterable[Dict[str, Any]]
+        self, fully_qualified_namespace: str, eventhub_name: str, consumer_group: str, **kwargs: Any
+    ) -> Iterable[Dict[str, Any]]:
         """List the updated checkpoints from the storage blob.
 
         :param str fully_qualified_namespace: The fully qualified namespace that the Event Hub belongs to.
@@ -351,5 +369,5 @@ def list_checkpoints(
             result.append(checkpoint)
         return result
 
-    def close(self):
+    def close(self) -> None:
         self._container_client.__exit__()

sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py

@@ -50,14 +50,20 @@
         "EventHubSharedKeyCredential",
         TokenCredential,
     ]
+    from ._consumer_client import EventHubConsumerClient
+    from ._producer_client import EventHubProducerClient
+    from ._transport._base import AmqpTransport
     try:
         from uamqp import Message as uamqp_Message
         from uamqp.authentication import JWTTokenAuth as uamqp_JWTTokenAuth
+        from uamqp import ReceiveClient as uamqp_AMQPRecieveClient
+        from uamqp import SendClient as uamqp_AMQPSendClient
     except ImportError:
-        uamqp_Message = None
-        uamqp_JWTTokenAuth = None
+        pass
     from ._pyamqp.message import Message
     from ._pyamqp.authentication import JWTTokenAuth
+    from ._pyamqp import ReceiveClient as pyamqp_AMQPRecieveClient
+    from ._pyamqp import SendClient as pyamqp_AMQPSendClient
 
 _LOGGER = logging.getLogger(__name__)
 _Address = collections.namedtuple("_Address", "hostname path")
@@ -333,7 +339,7 @@ def _from_connection_string(conn_str: str, **kwargs: Any) -> Dict[str, Any]:
             kwargs["credential"] = EventHubSharedKeyCredential(policy, key)
         return kwargs
 
-    def _create_auth(self) -> Union[uamqp_JWTTokenAuth, JWTTokenAuth]:
+    def _create_auth(self) -> Union["uamqp_JWTTokenAuth", JWTTokenAuth]:
         """
         Create an ~uamqp.authentication.SASTokenAuth instance
          to authenticate the session.
@@ -414,7 +420,8 @@ def _management_request(
                 mgmt_client.open(connection=conn)
                 while not mgmt_client.client_ready():
                     time.sleep(0.05)
-                mgmt_msg.application_properties[
+
+                cast(Dict[Union[str, bytes], Any], mgmt_msg.application_properties)[
                     "security_token"
                 ] = self._amqp_transport.get_updated_token(mgmt_auth)
                 status_code, description, response = self._amqp_transport.mgmt_client_request(
@@ -510,11 +517,22 @@ def _close(self) -> None:
         self._conn_manager.close_connection()
 
 
-class ConsumerProducerMixin(object):
-    def __enter__(self):
+class ConsumerProducerMixin():
+
+    def __init__(self) -> None:
+        self._handler: Union[
+            uamqp_AMQPRecieveClient,
+            pyamqp_AMQPRecieveClient,
+            uamqp_AMQPSendClient,
+            pyamqp_AMQPSendClient]
+        self._client: Union[EventHubConsumerClient, EventHubProducerClient]
+        self._amqp_transport: "AmqpTransport"
+        self._max_message_size_on_link: Optional[int] = None
+
+    def __enter__(self) -> ConsumerProducerMixin:
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         self.close()
 
     def _create_handler(self, auth):
@@ -526,8 +544,12 @@ def _check_closed(self):
                 f"{self._name} has been closed. Please create a new one to handle event data."
             )
 
-    def _open(self):
-        """Open the EventHubConsumer/EventHubProducer using the supplied connection."""
+    def _open(self) -> bool:
+        """Open the EventHubConsumer/EventHubProducer using the supplied connection.
+
+        :return: Whether the EventHubConsumer/EventHubProducer is ready to use.
+        :rtype: bool
+        """
         # pylint: disable=protected-access
         if not self.running:
             if self._handler:
@@ -544,9 +566,11 @@ def _open(self):
                 self._amqp_transport.get_remote_max_message_size(self._handler)
                 or self._amqp_transport.MAX_MESSAGE_LENGTH_BYTES
             )
-            self.running = True
+            self.running: bool = True
+            return True
+        return False
 
-    def _close_handler(self):
+    def _close_handler(self) -> None:
         if self._handler:
             self._handler.close()  # close the link (sharing connection) or connection (not sharing)
         self.running = False

sdk/eventhub/azure-eventhub/azure/eventhub/_common.py

@@ -64,8 +64,8 @@
             BatchMessage,
         )
     except ImportError:
-        Message = None
-        BatchMessage = None
+        pass
+
     from ._transport._base import AmqpTransport
 
 MessageContent = TypedDict("MessageContent", {"content": bytes, "content_type": str})
@@ -76,8 +76,8 @@
         bytes,
         bool,
         str,
-        Dict,
-        List,
+        Dict[str, Any],
+        List[Any],
         uuid.UUID,
     ]
 ]
@@ -348,12 +348,12 @@ def properties(self) -> Dict[Union[str, bytes], Any]:
         return self._raw_amqp_message.application_properties
 
     @properties.setter
-    def properties(self, value: Dict[Union[str, bytes], Any]):
+    def properties(self, value: Dict[Union[str, bytes], Any]) -> None:
         """Application-defined properties on the event.
 
         :param dict[str, any] or dict[bytes, any] value: The application properties for the EventData.
         """
-        properties = None if value is None else dict(value)
+        properties = None if value is None else value
         self._raw_amqp_message.application_properties = properties
 
     @property
@@ -550,7 +550,7 @@ def __init__(
         max_size_in_bytes: Optional[int] = None,
         partition_id: Optional[str] = None,
         partition_key: Optional[Union[str, bytes]] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         self._amqp_transport = kwargs.pop("amqp_transport", PyamqpTransport)
         self._tracing_attributes: Dict[str, Union[str, int]] = kwargs.pop("tracing_attributes", {})
@@ -571,7 +571,7 @@ def __init__(
             self._message, self._partition_key
         )
         self._size = self._amqp_transport.get_batch_message_encoded_size(self._message)
-        self.max_size_in_bytes = (
+        self.max_size_in_bytes: int = (
             max_size_in_bytes or self._amqp_transport.MAX_MESSAGE_LENGTH_BYTES
         )
 

sdk/eventhub/azure-eventhub/azure/eventhub/_connection_manager.py

@@ -49,7 +49,7 @@ class _ConnectionMode(Enum):
 
 
 class _SharedConnectionManager(object):  # pylint:disable=too-many-instance-attributes
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any):
         self._lock = Lock()
         self._conn: Union[Connection, uamqp_Connection] = None