[EventHubs] fix credential mypy (Azure#23548)

* fix mypy

* fix pylint

* fix sync

* Apply suggestions from code review

* apply review feedback

* move type checking to the front

Author: yunhaoling

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

@@ -42,8 +42,6 @@
     MGMT_STATUS_DESC,
 )
 
-if TYPE_CHECKING:
-    from azure.core.credentials import TokenCredential
 
 _LOGGER = logging.getLogger(__name__)
 _Address = collections.namedtuple("_Address", "hostname path")
@@ -264,9 +262,19 @@ def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
         return AccessToken(signature, expiry)
 
 
+if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
+    CredentialTypes = Union[
+        AzureSasCredential,
+        AzureNamedKeyCredential,
+        EventHubSharedKeyCredential,
+        TokenCredential
+    ]
+
+
 class ClientBase(object):  # pylint:disable=too-many-instance-attributes
     def __init__(self, fully_qualified_namespace, eventhub_name, credential, **kwargs):
-        # type: (str, str, Union[AzureSasCredential, TokenCredential, AzureNamedKeyCredential], Any) -> None
+        # type: (str, str, CredentialTypes, Any) -> None
         self.eventhub_name = eventhub_name
         if not eventhub_name:
             raise ValueError("The eventhub name can not be None or empty.")

sdk/eventhub/azure-eventhub/azure/eventhub/_consumer_client.py

@@ -15,11 +15,6 @@
 
 if TYPE_CHECKING:
     import datetime
-    from azure.core.credentials import (
-        TokenCredential,
-        AzureSasCredential,
-        AzureNamedKeyCredential,
-    )
     from typing import (  # pylint: disable=ungrouped-imports
         Any,
         Union,
@@ -31,6 +26,7 @@
     )
     from ._eventprocessor.partition_context import PartitionContext
     from ._common import EventData
+    from ._client_base import CredentialTypes
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -144,7 +140,7 @@ def __init__(
         fully_qualified_namespace,  # type: str
         eventhub_name,  # type: str
         consumer_group,  # type: str
-        credential,  # type: Union[AzureSasCredential, TokenCredential, AzureNamedKeyCredential]
+        credential,  # type: CredentialTypes
         **kwargs  # type: Any
     ):
         # type: (...) -> None

sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py

@@ -17,11 +17,7 @@
 from ._common import EventDataBatch, EventData
 
 if TYPE_CHECKING:
-    from azure.core.credentials import (
-        TokenCredential,
-        AzureSasCredential,
-        AzureNamedKeyCredential,
-    )
+    from ._client_base import CredentialTypes
 
 SendEventTypes = List[Union[EventData, AmqpAnnotatedMessage]]
 
@@ -92,7 +88,7 @@ def __init__(
         self,
         fully_qualified_namespace,  # type: str
         eventhub_name,  # type: str
-        credential,  # type: Union[AzureSasCredential, TokenCredential, AzureNamedKeyCredential]
+        credential,  # type: CredentialTypes
         **kwargs  # type: Any
     ):
         # type:(...) -> None

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_client_base_async.py

@@ -46,12 +46,70 @@
 
 if TYPE_CHECKING:
     from azure.core.credentials_async import AsyncTokenCredential
+    CredentialTypes = Union[
+        "EventHubSharedKeyCredential",
+        AsyncTokenCredential,
+        AzureSasCredential,
+        AzureNamedKeyCredential,
+    ]
 
     try:
         from typing_extensions import Protocol
     except ImportError:
         Protocol = object  # type: ignore
 
+    class AbstractConsumerProducer(Protocol):
+        @property
+        def _name(self):
+            # type: () -> str
+            """Name of the consumer or producer"""
+
+        @_name.setter
+        def _name(self, value):
+            pass
+
+        @property
+        def _client(self):
+            # type: () -> ClientBaseAsync
+            """The instance of EventHubComsumerClient or EventHubProducerClient"""
+
+        @_client.setter
+        def _client(self, value):
+            pass
+
+        @property
+        def _handler(self):
+            # type: () -> AMQPClientAsync
+            """The instance of SendClientAsync or ReceiveClientAsync"""
+
+        @property
+        def _internal_kwargs(self):
+            # type: () -> dict
+            """The dict with an event loop that users may pass in to wrap sync calls to async API.
+            It's furthur passed to uamqp APIs
+            """
+
+        @_internal_kwargs.setter
+        def _internal_kwargs(self, value):
+            pass
+
+        @property
+        def running(self):
+            # type: () -> bool
+            """Whether the consumer or producer is running"""
+
+        @running.setter
+        def running(self, value):
+            pass
+
+        def _create_handler(self, auth: authentication.JWTTokenAsync) -> None:
+            pass
+
+    _MIXIN_BASE = AbstractConsumerProducer
+else:
+    _MIXIN_BASE = object
+
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -144,9 +202,7 @@ def __init__(
         self,
         fully_qualified_namespace: str,
         eventhub_name: str,
-        credential: Union[
-            "AsyncTokenCredential", AzureSasCredential, AzureNamedKeyCredential
-        ],
+        credential: "CredentialTypes",
         **kwargs: Any
     ) -> None:
         self._internal_kwargs = get_dict_with_loop_if_needed(kwargs.get("loop", None))
@@ -376,60 +432,6 @@ async def _close_async(self) -> None:
         await self._conn_manager_async.close_connection()
 
 
-if TYPE_CHECKING:
-
-    class AbstractConsumerProducer(Protocol):
-        @property
-        def _name(self):
-            # type: () -> str
-            """Name of the consumer or producer"""
-
-        @_name.setter
-        def _name(self, value):
-            pass
-
-        @property
-        def _client(self):
-            # type: () -> ClientBaseAsync
-            """The instance of EventHubComsumerClient or EventHubProducerClient"""
-
-        @_client.setter
-        def _client(self, value):
-            pass
-
-        @property
-        def _handler(self):
-            # type: () -> AMQPClientAsync
-            """The instance of SendClientAsync or ReceiveClientAsync"""
-
-        @property
-        def _internal_kwargs(self):
-            # type: () -> dict
-            """The dict with an event loop that users may pass in to wrap sync calls to async API.
-            It's furthur passed to uamqp APIs
-            """
-
-        @_internal_kwargs.setter
-        def _internal_kwargs(self, value):
-            pass
-
-        @property
-        def running(self):
-            # type: () -> bool
-            """Whether the consumer or producer is running"""
-
-        @running.setter
-        def running(self, value):
-            pass
-
-        def _create_handler(self, auth: authentication.JWTTokenAsync) -> None:
-            pass
-
-    _MIXIN_BASE = AbstractConsumerProducer
-else:
-    _MIXIN_BASE = object
-
-
 class ConsumerProducerMixin(_MIXIN_BASE):
     async def __aenter__(self):
         return self

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_client_async.py

@@ -18,8 +18,6 @@
     Awaitable,
 )
 
-from azure.core.credentials import AzureSasCredential, AzureNamedKeyCredential
-
 from ._eventprocessor.event_processor import EventProcessor
 from ._consumer_async import EventHubConsumer
 from ._client_base_async import ClientBaseAsync
@@ -28,7 +26,7 @@
 
 
 if TYPE_CHECKING:
-    from azure.core.credentials_async import AsyncTokenCredential
+    from ._client_base_async import CredentialTypes
     from uamqp.constants import TransportType
     from ._eventprocessor.partition_context import PartitionContext
     from ._eventprocessor.checkpoint_store import CheckpointStore
@@ -147,7 +145,7 @@ def __init__(
         fully_qualified_namespace: str,
         eventhub_name: str,
         consumer_group: str,
-        credential: Union["AsyncTokenCredential", AzureSasCredential, AzureNamedKeyCredential],
+        credential: "CredentialTypes",
         **kwargs
     ) -> None:
         self._checkpoint_store = kwargs.pop("checkpoint_store", None)

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py

@@ -8,8 +8,6 @@
 from typing import Any, Union, TYPE_CHECKING, List, Optional, Dict, cast
 from uamqp import constants
 
-from azure.core.credentials import AzureSasCredential, AzureNamedKeyCredential
-
 from ..exceptions import ConnectError, EventHubError
 from ..amqp import AmqpAnnotatedMessage
 from ._client_base_async import ClientBaseAsync
@@ -18,7 +16,7 @@
 from .._common import EventDataBatch, EventData
 
 if TYPE_CHECKING:
-    from azure.core.credentials_async import AsyncTokenCredential
+    from ._client_base_async import CredentialTypes
     from uamqp.constants import TransportType  # pylint: disable=ungrouped-imports
 
 SendEventTypes = List[Union[EventData, AmqpAnnotatedMessage]]
@@ -90,9 +88,7 @@ def __init__(
         self,
         fully_qualified_namespace: str,
         eventhub_name: str,
-        credential: Union[
-            "AsyncTokenCredential", AzureSasCredential, AzureNamedKeyCredential
-        ],
+        credential: "CredentialTypes",
         **kwargs
     ) -> None:
         super(EventHubProducerClient, self).__init__(