[SB] Sync Azure AMQP Typing changes (#35472)

* bring over typing update

* fix type of node

Author: kashifkhan

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_client_async.py

@@ -4,12 +4,13 @@
 # license information.
 #--------------------------------------------------------------------------
 # TODO: Check types of kwargs (issue exists for this)
+# pylint: disable=too-many-lines
 import asyncio
 import logging
 import time
 import queue
 from functools import partial
-from typing import Any, Dict, Optional, Tuple, Union, overload, cast
+from typing import Any, Callable, Coroutine, List, Dict, Optional, Tuple, Union, overload, cast
 from typing_extensions import Literal
 import certifi
 
@@ -23,7 +24,7 @@
     SendClient as SendClientSync,
     Outcomes
 )
-from ..message import _MessageDelivery
+from ..message import _MessageDelivery, Message
 from ..constants import (
     MessageDeliveryState,
     SEND_DISPOSITION_ACCEPT,
@@ -350,7 +351,16 @@ async def do_work_async(self, **kwargs):
             return True
         return await self._client_run_async(**kwargs)
 
-    async def mgmt_request_async(self, message, **kwargs):
+    async def mgmt_request_async(
+            self,
+            message,
+            *,
+            operation: Optional[Union[bytes, str]] = None,
+            operation_type: Optional[Union[bytes, str]] = None,
+            node: str = "$management",
+            timeout: float = 0,
+            **kwargs
+        ):
         """
         :param message: The message to send in the management request.
         :type message: ~pyamqp.message.Message
@@ -370,10 +380,6 @@ async def mgmt_request_async(self, message, **kwargs):
         # The method also takes "status_code_field" and "status_description_field"
         # keyword arguments as alternate names for the status code and description
         # in the response body. Those two keyword arguments are used in Azure services only.
-        operation = kwargs.pop("operation", None)
-        operation_type = kwargs.pop("operation_type", None)
-        node = kwargs.pop("node", "$management")
-        timeout = kwargs.pop('timeout', 0)
         async with self._mgmt_link_lock_async:
             try:
                 mgmt_link = self._mgmt_links[node]
@@ -569,8 +575,7 @@ async def _on_send_complete_async(self, message_delivery, reason, state):
                 condition=ErrorCondition.UnknownError
             )
 
-    async def _send_message_impl_async(self, message, **kwargs):
-        timeout = kwargs.pop("timeout", 0)
+    async def _send_message_impl_async(self, message, *, timeout: float = 0):
         expire_time = (time.time() + timeout) if timeout else None
         await self.open_async()
         message_delivery = _MessageDelivery(
@@ -604,11 +609,19 @@ async def _send_message_impl_async(self, message, **kwargs):
                 # This is a default handler
                 raise MessageException(condition=ErrorCondition.UnknownError, description="Send failed.") from None
 
-    async def send_message_async(self, message, **kwargs):
+    async def send_message_async(self, message, *, timeout: float = 0, **kwargs):
         """
         :param ~pyamqp.message.Message message: The message to send.
+        :keyword float timeout: timeout in seconds. If set to
+         0, the client will continue to wait until the message is sent or error happens. The
+         default is 0.
         """
-        await self._do_retryable_operation_async(self._send_message_impl_async, message=message, **kwargs)
+        await self._do_retryable_operation_async(
+            self._send_message_impl_async,
+            message=message,
+            timeout=timeout,
+            **kwargs
+        )
 
 
 class ReceiveClientAsync(ReceiveClientSync, AMQPClientAsync):
@@ -767,12 +780,17 @@ async def _message_received_async(self, frame, message):
         if not self._streaming_receive:
             self._received_messages.put((frame, message))
 
-    async def _receive_message_batch_impl_async(self, max_batch_size=None, on_message_received=None, timeout=0):
+    async def _receive_message_batch_impl_async(
+            self,
+            max_batch_size: Optional[int] = None,
+            on_message_received: Optional[Callable] = None,
+            timeout: float = 0,
+        ):
         self._message_received_callback = on_message_received
         max_batch_size = max_batch_size or self._link_credit
         timeout_time = time.time() + timeout if timeout else 0
         receiving = True
-        batch = []
+        batch: List[Message] = []
         await self.open_async()
         while len(batch) < max_batch_size:
             try:
@@ -825,7 +843,14 @@ async def close_async(self):
         self._received_messages = queue.Queue()
         await super(ReceiveClientAsync, self).close_async()
 
-    async def receive_message_batch_async(self, **kwargs):
+    async def receive_message_batch_async( # pylint: disable=unused-argument
+            self,
+            *,
+            max_batch_size: Optional[int] = None,
+            on_message_received: Optional[Callable] = None,
+            timeout: float = 0,
+            **kwargs
+        ) -> Coroutine[Any, Any, list]:
         """Receive a batch of messages. Messages returned in the batch have already been
         accepted - if you wish to add logic to accept or reject messages based on custom
         criteria, pass in a callback. This method will return as soon as some messages are
@@ -844,12 +869,14 @@ async def receive_message_batch_async(self, **kwargs):
          0, the client will continue to wait until at least one message is received. The
          default is 0.
         :paramtype timeout: float
-        :return: Retryable operation coroutine.
-        :rtype: Coroutine[Any, Any, Any]
+        :returns: Retryable operation coroutine.
+        :rtype: Coroutine[Any, Any, list]
         """
         return await self._do_retryable_operation_async(
             self._receive_message_batch_impl_async,
-            **kwargs
+            max_batch_size=max_batch_size,
+            on_message_received=on_message_received,
+            timeout=timeout,
         )
 
     async def receive_messages_iter_async(self, timeout=None, on_message_received=None):

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_management_link_async.py

@@ -7,6 +7,7 @@
 import time
 import logging
 from functools import partial
+from typing import Optional, Union
 
 from ..management_link import PendingManagementOperation
 from ._sender_async import SenderLink
@@ -189,7 +190,16 @@ async def open(self):
         await self._response_link.attach()
         await self._request_link.attach()
 
-    async def execute_operation(self, message, on_execute_operation_complete, **kwargs):
+    async def execute_operation(
+            self,
+            message,
+            on_execute_operation_complete,
+            *,
+            operation: Optional[Union[bytes, str]] = None,
+            type: Optional[Union[bytes, str]] = None,
+            locales: Optional[str] = None,
+            timeout: Optional[float] = None
+        ):
         """Execute a request and wait on a response.
 
         :param message: The message to send in the management request.
@@ -212,11 +222,10 @@ async def execute_operation(self, message, on_execute_operation_complete, **kwar
          to the management request must be received.
         :rtype: None
         """
-        timeout = kwargs.get("timeout")
-        message.application_properties["operation"] = kwargs.get("operation")
-        message.application_properties["type"] = kwargs.get("type")
-        if "locales" in kwargs:
-            message.application_properties["locales"] = kwargs.get("locales")
+        message.application_properties["operation"] = operation
+        message.application_properties["type"] = type
+        if locales:
+            message.application_properties["locales"] = locales
         try:
             # TODO: namedtuple is immutable, which may push us to re-think about the namedtuple approach for Message
             new_properties = message.properties._replace(message_id=self.next_message_id)

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/client.py

@@ -15,12 +15,12 @@
 import time
 import uuid
 from functools import partial
-from typing import Any, Dict, Optional, Tuple, Union, overload, cast
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, overload, cast
 import certifi
 from typing_extensions import Literal
 
 from ._connection import Connection
-from .message import _MessageDelivery
+from .message import _MessageDelivery, Message
 from .error import (
     AMQPException,
     ErrorCondition,
@@ -426,7 +426,16 @@ def do_work(self, **kwargs):
             return True
         return self._client_run(**kwargs)
 
-    def mgmt_request(self, message, **kwargs):
+    def mgmt_request(
+            self,
+            message,
+            *,
+            operation: Optional[Union[str, bytes]] = None,
+            operation_type: Optional[Union[str, bytes]] = None,
+            node: str = "$management",
+            timeout: float = 0,
+            **kwargs
+        ):
         """
         :param message: The message to send in the management request.
         :type message: ~pyamqp.message.Message
@@ -446,10 +455,6 @@ def mgmt_request(self, message, **kwargs):
         # The method also takes "status_code_field" and "status_description_field"
         # keyword arguments as alternate names for the status code and description
         # in the response body. Those two keyword arguments are used in Azure services only.
-        operation = kwargs.pop("operation", None)
-        operation_type = kwargs.pop("operation_type", None)
-        node = kwargs.pop("node", "$management")
-        timeout = kwargs.pop("timeout", 0)
         with self._mgmt_link_lock:
             try:
                 mgmt_link = self._mgmt_links[node]
@@ -659,8 +664,7 @@ def _on_send_complete(self, message_delivery, reason, state):
                 message_delivery, condition=ErrorCondition.UnknownError
             )
 
-    def _send_message_impl(self, message, **kwargs):
-        timeout = kwargs.pop("timeout", 0)
+    def _send_message_impl(self, message, *, timeout: float = 0):
         expire_time = (time.time() + timeout) if timeout else None
         self.open()
         message_delivery = _MessageDelivery(
@@ -692,14 +696,19 @@ def _send_message_impl(self, message, **kwargs):
                     condition=ErrorCondition.UnknownError, description="Send failed."
                 ) from None
 
-    def send_message(self, message, **kwargs):
+    def send_message(self, message, *, timeout: float = 0, **kwargs):
         """
         :param ~pyamqp.message.Message message:
         :keyword float timeout: timeout in seconds. If set to
          0, the client will continue to wait until the message is sent or error happens. The
          default is 0.
         """
-        self._do_retryable_operation(self._send_message_impl, message=message, **kwargs)
+        self._do_retryable_operation(
+            self._send_message_impl,
+            message=message,
+            timeout=timeout,
+            **kwargs
+        )
 
 
 class ReceiveClient(AMQPClient): # pylint:disable=too-many-instance-attributes
@@ -877,13 +886,16 @@ def _message_received(self, frame, message):
             self._received_messages.put((frame, message))
 
     def _receive_message_batch_impl(
-        self, max_batch_size=None, on_message_received=None, timeout=0
+        self,
+        max_batch_size: Optional[int] = None,
+        on_message_received: Optional[Callable] = None,
+        timeout: float = 0,
     ):
         self._message_received_callback = on_message_received
         max_batch_size = max_batch_size or self._link_credit
         timeout = time.time() + timeout if timeout else 0
         receiving = True
-        batch = []
+        batch: List[Message] = []
         self.open()
         while len(batch) < max_batch_size:
             try:
@@ -928,7 +940,14 @@ def close(self):
         self._received_messages = queue.Queue()
         super(ReceiveClient, self).close()
 
-    def receive_message_batch(self, **kwargs):
+    def receive_message_batch( # pylint: disable=unused-argument
+            self,
+            *,
+            max_batch_size: Optional[int] = None,
+            on_message_received: Optional[Callable] = None,
+            timeout: float = 0,
+            **kwargs
+        ):
         """Receive a batch of messages. Messages returned in the batch have already been
         accepted - if you wish to add logic to accept or reject messages based on custom
         criteria, pass in a callback. This method will return as soon as some messages are
@@ -950,7 +969,12 @@ def receive_message_batch(self, **kwargs):
         :return: A list of messages.
         :rtype: list[~pyamqp.message.Message]
         """
-        return self._do_retryable_operation(self._receive_message_batch_impl, **kwargs)
+        return self._do_retryable_operation(
+            self._receive_message_batch_impl,
+            max_batch_size=max_batch_size,
+            on_message_received=on_message_received,
+            timeout=timeout
+        )
 
     def receive_messages_iter(self, timeout=None, on_message_received=None):
         """Receive messages by generator. Messages returned in the generator have already been

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/management_link.py

@@ -8,6 +8,7 @@
 import logging
 from functools import partial
 from collections import namedtuple
+from typing import Optional, Union
 
 from .sender import SenderLink
 from .receiver import ReceiverLink
@@ -193,7 +194,11 @@ def execute_operation(
         self,
         message,
         on_execute_operation_complete,
-        **kwargs
+        *,
+        operation: Optional[Union[bytes, str]] = None,
+        type: Optional[Union[bytes, str]] = None,
+        locales: Optional[str] = None,
+        timeout: Optional[float] = None
     ):
         """Execute a request and wait on a response.
 
@@ -217,11 +222,10 @@ def execute_operation(
          to the management request must be received.
         :rtype: None
         """
-        timeout = kwargs.get("timeout")
-        message.application_properties["operation"] = kwargs.get("operation")
-        message.application_properties["type"] = kwargs.get("type")
-        if "locales" in kwargs:
-            message.application_properties["locales"] = kwargs.get("locales")
+        message.application_properties["operation"] = operation
+        message.application_properties["type"] = type
+        if locales:
+            message.application_properties["locales"] = locales
         try:
             # TODO: namedtuple is immutable, which may push us to re-think about the namedtuple approach for Message
             new_properties = message.properties._replace(message_id=self.next_message_id)

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py

@@ -926,7 +926,7 @@ def mgmt_client_request(
         *,
         operation: bytes,
         operation_type: bytes,
-        node: bytes,
+        node: str,
         timeout: int,
         callback: Callable
     ) -> "ServiceBusReceivedMessage":
@@ -936,7 +936,7 @@ def mgmt_client_request(
         :param ~pyamqp.message.Message mgmt_msg: Message.
         :keyword bytes operation: Operation.
         :keyword bytes operation_type: Op type.
-        :keyword bytes node: Mgmt target.
+        :keyword str node: Mgmt target.
         :keyword int timeout: Timeout.
         :keyword callable callback: Callback to process request response.
         :return: ServiceBusReceivedMessage

sdk/servicebus/azure-servicebus/azure/servicebus/aio/_transport/_pyamqp_transport_async.py

@@ -418,7 +418,7 @@ async def mgmt_client_request_async(
         *,
         operation: bytes,
         operation_type: bytes,
-        node: bytes,
+        node: str,
         timeout: int,
         callback: Callable
     ) -> "ServiceBusReceivedMessage":
@@ -428,7 +428,7 @@ async def mgmt_client_request_async(
         :param ~pyamqp.message.Message mgmt_msg: Message.
         :keyword bytes operation: Operation.
         :keyword bytes operation_type: Op type.
-        :keyword bytes node: Mgmt target.
+        :keyword str node: Mgmt target.
         :keyword int timeout: Timeout.
         :keyword callable callback: Callback to process request response.
         :return: The result returned by the mgmt request.