Azure · Apr 1, 2023 · Jan 27, 2023 · Jan 30, 2023 · Jan 30, 2023 · Jan 30, 2023
@@ -947,8 +947,17 @@ def __repr__(self) -> str:  # pylint: disable=too-many-branches,too-many-stateme
             message_repr += ", locked_until_utc=<read-error>"
         return "ServiceBusReceivedMessage({})".format(message_repr)[:1024]
 
-    @property
+    @property # type: ignore[misc]  # TODO: ignoring error to copy over setter, since it's inherited
     def message(self) -> LegacyMessage:
+        """DEPRECATED: Get the underlying LegacyMessage.
+         This is deprecated and will be removed in a later release.
+
+        :rtype: LegacyMessage
+        """
+        warnings.warn(
+            "The `message` property is deprecated and will be removed in future versions.",
+            DeprecationWarning,
+        )
         if not self._uamqp_message:
             if not self._settled:
                 settler = self._receiver._handler # pylint:disable=protected-access

@@ -4,7 +4,7 @@
 # license information.
 # -------------------------------------------------------------------------
 import uuid
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 from ..exceptions import MessageAlreadySettled
 from .constants import (
@@ -15,10 +15,11 @@
 
 if TYPE_CHECKING:
     from .._transport._base import AmqpTransport
+    from ..aio._transport._base_async import AmqpTransportAsync
 
 class ReceiverMixin(object):  # pylint: disable=too-many-instance-attributes
     def _populate_attributes(self, **kwargs):
-        self._amqp_transport: "AmqpTransport"
+        self._amqp_transport: Union["AmqpTransport", "AmqpTransportAsync"]
         if kwargs.get("subscription_name"):
             self._subscription_name = kwargs.get("subscription_name")
             self._is_subscription = True

@@ -75,6 +75,7 @@
     from .receiver_mixins import ReceiverMixin
     from .._servicebus_session import BaseSession
     from .._transport._base import AmqpTransport
+    from ..aio._transport._base_async import AmqpTransportAsync
 
     MessagesType = Union[
         Mapping[str, Any],
@@ -304,7 +305,7 @@ def receive_trace_context_manager(
 
 def trace_message(
     message: Union["uamqp_Message", "pyamqp_Message"],
-    amqp_transport: "AmqpTransport",
+    amqp_transport: Union["AmqpTransport", "AmqpTransportAsync"],
     parent_span: Optional["AbstractSpan"] = None
 ) -> Union["uamqp_Message", "pyamqp_Message"]:
     """Add tracing information to this message.

@@ -142,8 +142,9 @@ class ApacheFilters(object):
         'default_outcome',
         'outcomes',
         'capabilities'
-    ])
-Source.__new__.__defaults__ = (None,) * len(Source._fields) # type: ignore
+    ],
+    defaults=(None,) * 11 # type: ignore
+    )
 Source._code = 0x00000028 # type: ignore # pylint: disable=protected-access
 Source._definition = ( # type: ignore # pylint: disable=protected-access
     FIELD("address", AMQPTypes.string, False, None, False),
@@ -227,9 +228,10 @@ class ApacheFilters(object):
         'dynamic',
         'dynamic_node_properties',
         'capabilities'
-    ])
+    ],
+    defaults=(None,) * 7 # type: ignore
+    )
 Target._code = 0x00000029 # type: ignore # pylint: disable=protected-access
-Target.__new__.__defaults__ = (None,) * len(Target._fields) # type: ignore # type: ignore # pylint: disable=protected-access
 Target._definition = ( # type: ignore # pylint: disable=protected-access
     FIELD("address", AMQPTypes.string, False, None, False),
     FIELD("durable", AMQPTypes.uint, False, "none", False),

@@ -20,9 +20,10 @@
         'ttl',
         'first_acquirer',
         'delivery_count'
-    ])
+    ],
+    defaults=(None,) * 5 # type: ignore
+    )
 Header._code = 0x00000070 # type: ignore # pylint:disable=protected-access
-Header.__new__.__defaults__ = (None,) * len(Header._fields) # type: ignore
 Header._definition = ( # type: ignore # pylint:disable=protected-access
     FIELD("durable", AMQPTypes.boolean, False, None, False),
     FIELD("priority", AMQPTypes.ubyte, False, None, False),
@@ -91,7 +92,9 @@
         'group_id',
         'group_sequence',
         'reply_to_group_id'
-    ])
+    ],
+    defaults=(None,) * 13 # type: ignore
+    )
 Properties._code = 0x00000073 # type: ignore # pylint:disable=protected-access
 Properties.__new__.__defaults__ = (None,) * len(Properties._fields) # type: ignore
 Properties._definition = ( # type: ignore # pylint:disable=protected-access
@@ -177,8 +180,9 @@
         'sequence',
         'value',
         'footer',
-    ])
-Message.__new__.__defaults__ = (None,) * len(Message._fields) # type: ignore
+    ],
+    defaults=(None,) * 9 # type: ignore
+    )
 Message._code = 0 # type: ignore # pylint:disable=protected-access
 Message._definition = ( # type: ignore # pylint:disable=protected-access
     (0x00000070, FIELD("header", Header, False, None, False)),

@@ -116,7 +116,7 @@ def __init__(
         **kwargs: Any
     ) -> None:
         uamqp_transport = kwargs.pop("uamqp_transport", False)
-        if uamqp_transport and not UamqpTransport:
+        if uamqp_transport and UamqpTransport is None:
             raise ValueError("To use the uAMQP transport, please install `uamqp>=1.6.3,<2.0.0`.")
         self._amqp_transport = UamqpTransport if uamqp_transport else PyamqpTransport
 

@@ -340,7 +340,7 @@ def _create_handler(self, auth: Union["pyamqp_JWTTokenAuth", "uamqp_JWTTokenAuth
         if self._prefetch_count == 1:
             # pylint: disable=protected-access
             self._handler._message_received = functools.partial(
-                self._amqp_transport.enhanced_message_received,
+                self._amqp_transport.enhanced_message_received, # type: ignore[attr-defined]
                 self
             )
 

@@ -409,11 +409,12 @@ def send_messages(
             raise ValueError("The timeout must be greater than 0.")
 
         try:  # Short circuit noop if an empty list or batch is provided.
-            if len(message) == 0:  # pylint: disable=len-as-condition
+            if len(cast(Union[List, ServiceBusMessageBatch], message)) == 0:  # pylint: disable=len-as-condition
                 return
         except TypeError:   # continue if ServiceBusMessage
             pass
 
+        obj_message: Union[ServiceBusMessage, ServiceBusMessageBatch]
         with send_trace_context_manager() as send_span:
             if isinstance(message, ServiceBusMessageBatch):
                 # If AmqpTransports are not the same, create batch with correct BatchMessage.
@@ -423,7 +424,7 @@ def send_messages(
                     batch._from_list(message._messages, send_span)  # type: ignore
                     obj_message = batch
                 else:
-                    obj_message: "MessageObjTypes" = message
+                    obj_message = message
             else:
                 obj_message = transform_outbound_messages(  # type: ignore
                     message, ServiceBusMessage, self._amqp_transport.to_outgoing_amqp_message
@@ -439,7 +440,6 @@ def send_messages(
                         parent_span=send_span
                     )
 
-            obj_message = cast(Union[ServiceBusMessage, ServiceBusMessageBatch], obj_message)
             if send_span:
                 self._add_span_request_attributes(send_span)
             self._do_retryable_operation(

@@ -3,7 +3,7 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 from __future__ import annotations
-from typing import Union, TYPE_CHECKING, Dict, Optional, List, Any, Callable
+from typing import Union, TYPE_CHECKING, Dict, Any, Callable
 from abc import ABC, abstractmethod
 
 if TYPE_CHECKING:
@@ -22,18 +22,20 @@ class AmqpTransport(ABC):   # pylint: disable=too-many-public-methods
     MAX_FRAME_SIZE_BYTES: int
     MAX_MESSAGE_LENGTH_BYTES: int
     TIMEOUT_FACTOR: int
-    #CONNECTION_CLOSING_STATES: Tuple
     TRANSPORT_IDENTIFIER: str
 
+    ServiceBusToAMQPReceiveModeMap: Dict[str, Any]
+
     # define symbols
     PRODUCT_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes]
     VERSION_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes]
     FRAMEWORK_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes]
     PLATFORM_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes]
     USER_AGENT_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes]
     PROP_PARTITION_KEY_AMQP_SYMBOL: Union[uamqp_types.AMQPSymbol, str, bytes]
-    AMQP_LONG_VALUE: Union[uamqp_types.AMQPLong, Dict[str, str]]
-    AMQP_ARRAY_VALUE: Union[uamqp_types.AMQPArray, Dict[str, str]]
+    AMQP_LONG_VALUE: Callable
+    AMQP_ARRAY_VALUE: Callable
+    AMQP_UINT_VALUE: Callable
 
     @staticmethod
     @abstractmethod
@@ -45,7 +47,7 @@ def build_message(**kwargs):
 
     @staticmethod
     @abstractmethod
-    def build_batch_message(**kwargs):
+    def build_batch_message(data):
         """
         Creates a uamqp.BatchMessage or pyamqp.BatchMessage with given arguments.
         :rtype: uamqp.BatchMessage or pyamqp.BatchMessage

@@ -7,7 +7,7 @@
 import time
 import datetime
 from datetime import timezone
-from typing import Optional, Tuple, cast, List, TYPE_CHECKING, Any, Callable, Dict
+from typing import Optional, Tuple, cast, List, TYPE_CHECKING, Any, Callable, Dict, Union, Iterator, Type
 
 from .._pyamqp import (
     utils,
@@ -93,7 +93,7 @@
 
 if TYPE_CHECKING:
     from logging import Logger
-    from ..amqp import AmqpAnnotatedMessage
+    from ..amqp import AmqpAnnotatedMessage, AmqpMessageHeader, AmqpMessageProperties
     from .._servicebus_receiver import ServiceBusReceiver
     from .._servicebus_sender import ServiceBusSender
     from .._common.message import ServiceBusReceivedMessage, ServiceBusMessage, ServiceBusMessageBatch
@@ -183,9 +183,9 @@ class PyamqpTransport(AmqpTransport):   # pylint: disable=too-many-public-method
     #ERROR_CONDITIONS = [condition.value for condition in ErrorCondition]
 
     # amqp value types
-    AMQP_LONG_VALUE = amqp_long_value
-    AMQP_ARRAY_VALUE = amqp_array_value
-    AMQP_UINT_VALUE = amqp_uint_value
+    AMQP_LONG_VALUE: Callable = amqp_long_value
+    AMQP_ARRAY_VALUE: Callable = amqp_array_value
+    AMQP_UINT_VALUE: Callable = amqp_uint_value
 
     # errors
     TIMEOUT_ERROR = TimeoutError
@@ -243,7 +243,8 @@ def to_outgoing_amqp_message(annotated_message: "AmqpAnnotatedMessage") -> "Mess
         ttl_set = False
         header_vals = annotated_message.header.values() if annotated_message.header else None
         # If header and non-None header values, create outgoing header.
-        if annotated_message.header and header_vals.count(None) != len(header_vals):
+        if header_vals and header_vals.count(None) != len(header_vals):
+            annotated_message.header = cast("AmqpMessageHeader", annotated_message.header)
             message_header = Header(
                 delivery_count=annotated_message.header.delivery_count,
                 ttl=annotated_message.header.time_to_live,
@@ -264,7 +265,8 @@ def to_outgoing_amqp_message(annotated_message: "AmqpAnnotatedMessage") -> "Mess
         message_properties = None
         properties_vals = annotated_message.properties.values() if annotated_message.properties else None
         # If properties and non-None properties values, create outgoing properties.
-        if annotated_message.properties and properties_vals.count(None) != len(properties_vals):
+        if properties_vals and properties_vals.count(None) != len(properties_vals):
+            annotated_message.properties = cast("AmqpMessageProperties", annotated_message.properties)
             creation_time = None
             absolute_expiry_time = None
             if ttl_set:
@@ -292,7 +294,7 @@ def to_outgoing_amqp_message(annotated_message: "AmqpAnnotatedMessage") -> "Mess
                 reply_to_group_id=annotated_message.properties.reply_to_group_id,
             )
         elif ttl_set:
-            message_properties = Properties(
+            message_properties = Properties(    # type: ignore[call-arg]
                 creation_time=creation_time_from_ttl if ttl_set else None,
                 absolute_expiry_time=absolute_expiry_time_from_ttl if ttl_set else None,
             )
@@ -457,7 +459,7 @@ def create_send_client(
     @staticmethod
     def send_messages(
         sender: "ServiceBusSender",
-        message: "Message",
+        message: Union["ServiceBusMessage", "ServiceBusMessageBatch"],
         logger: "Logger",
         timeout: int,
         last_exception: Optional[Exception]
@@ -507,7 +509,7 @@ def create_source(source: "Source", session_filter: Optional[str]) -> "Source":
         :param bytes session_filter: Required.
         """
         filter_map = {SESSION_FILTER: session_filter}
-        source = Source(address=source, filters=filter_map)
+        source = Source(address=source, filters=filter_map) # type: ignore[call-arg]
         return source
 
     @staticmethod
@@ -584,7 +586,7 @@ def on_attach(
     @staticmethod
     def iter_contextual_wrapper(
         receiver: "ServiceBusReceiver", max_wait_time: Optional[int] = None
-    ) -> "ServiceBusReceivedMessage":
+    ) -> Iterator["ServiceBusReceivedMessage"]:
         """The purpose of this wrapper is to allow both state restoration (for multiple concurrent iteration)
         and per-iter argument passing that requires the former."""
         while True:
@@ -609,7 +611,9 @@ def iter_next(
             receiver._open()
             if not receiver._message_iter or wait_time:
                 receiver._message_iter = receiver._handler.receive_messages_iter(timeout=wait_time)
-            pyamqp_message = next(receiver._message_iter)
+            pyamqp_message = next(
+                cast(Iterator["Message"], receiver._message_iter)
+            )
             message = receiver._build_received_message(pyamqp_message)
             if (
                 receiver._auto_lock_renewer
@@ -640,7 +644,7 @@ def enhanced_message_received(
     @staticmethod
     def build_received_message(
         receiver: "ServiceBusReceiver",
-        message_type: "ServiceBusReceivedMessage",
+        message_type: Type["ServiceBusReceivedMessage"],
         received: "Message"
     ) -> "ServiceBusReceivedMessage":
         """
@@ -656,7 +660,7 @@ def build_received_message(
     @staticmethod
     def get_current_time(
         handler: "ReceiveClient"
-    ) -> int:  # pylint: disable=unused-argument
+    ) -> float:  # pylint: disable=unused-argument
         """
         Gets the current time.
         """
@@ -717,7 +721,9 @@ def settle_message_via_receiver_link(
 
     @staticmethod
     def parse_received_message(
-        message: "Message", message_type: "ServiceBusReceivedMessage", **kwargs: Any
+        message: "Message",
+        message_type: Type["ServiceBusReceivedMessage"],
+        **kwargs: Any
     ) -> List["ServiceBusReceivedMessage"]:
         """
         Parses peek/deferred op messages into ServiceBusReceivedMessage.
@@ -836,10 +842,10 @@ def mgmt_client_request(
     @staticmethod
     def _handle_amqp_exception_with_condition(
         logger: "Logger",
-        condition: "ErrorCondition",
+        condition: Optional["ErrorCondition"],
         description: str,
-        exception: "AMQPException" = None,
-        status_code: str = None
+        exception: Optional["AMQPException"] = None,
+        status_code: Optional[str] = None
     ) -> "ServiceBusError":
         # handling AMQP Errors that have the condition field or the mgmt handler
         logger.info(
@@ -848,6 +854,7 @@ def _handle_amqp_exception_with_condition(
             condition,
             description,
         )
+        error_cls: Type["ServiceBusError"]
         if isinstance(exception, AuthenticationException):
             logger.info("AMQP Connection authentication error occurred: (%r).", exception)
             error_cls = ServiceBusAuthenticationError
@@ -870,7 +877,7 @@ def _handle_amqp_exception_with_condition(
         elif condition == ErrorCondition.UnknownError or isinstance(exception, AMQPConnectionError):
             error_cls = ServiceBusConnectionError
         else:
-            error_cls = _ERROR_CODE_TO_ERROR_MAPPING.get(condition, ServiceBusError)
+            error_cls = _ERROR_CODE_TO_ERROR_MAPPING.get(cast(bytes, condition), ServiceBusError)
 
         error = error_cls(
             message=description,
@@ -889,9 +896,9 @@ def _handle_amqp_exception_with_condition(
     def handle_amqp_mgmt_error(
         logger: "Logger",
         error_description: "str",
-        condition: "ErrorCondition" = None,
-        description: str = None,
-        status_code: str = None
+        condition: Optional["ErrorCondition"] = None,
+        description: Optional[str] = None,
+        status_code: Optional[str] = None
     ) -> "ServiceBusError":
         if description:
             error_description += f" {description}."