Support more types in publish payload

Author: mvollrath

ohmqtt/client.py

@@ -122,7 +122,7 @@ def shutdown(self) -> None:
     def publish(
         self,
         topic: str,
-        payload: bytes,
+        payload: bytes | bytearray | str,
         *,
         qos: int | MQTTQoS = MQTTQoS.Q0,
         retain: bool = False,
@@ -132,12 +132,16 @@ def publish(
         """Publish a message to a topic.
 
         :param topic: The topic to publish to.
-        :param payload: The payload of the message.
+        :param payload: The payload of the message. If a string is provided, it will be encoded as UTF-8.
         :param qos: The QoS level for the message (0, 1, or 2).
         :param retain: If True, the message will be retained by the broker.
         :param properties: Properties for the PUBLISH packet.
         :param alias_policy: The policy for using automatic topic aliases.
         """
+        if isinstance(payload, str):
+            payload = payload.encode("utf-8")
+        elif not isinstance(payload, bytes):
+            payload = bytes(payload)
         if not isinstance(qos, MQTTQoS):
             qos = MQTTQoS(qos)
         properties = properties if properties is not None else None

tests/test_client.py

@@ -77,15 +77,16 @@ def test_client_shutdown(mocker: MockerFixture, mock_connection: Mock, mock_hand
     mock_connection.shutdown.assert_called_once()
 
 
+@pytest.mark.parametrize("payload", [b"test_payload", bytearray(b"test_payload"), "test_payload"])
 @pytest.mark.parametrize("qos", [0, 1, 2, MQTTQoS.Q0, MQTTQoS.Q1, MQTTQoS.Q2])
-def test_client_publish(qos: int | MQTTQoS, mocker: MockerFixture, mock_connection: Mock, mock_handlers: MagicMock,
-                       mock_session: Mock, mock_subscriptions: Mock) -> None:
+def test_client_publish(payload: bytes | bytearray | str, qos: int | MQTTQoS, mocker: MockerFixture,
+                        mock_connection: Mock, mock_handlers: MagicMock, mock_session: Mock, mock_subscriptions: Mock) -> None:
     client = Client()
 
     mock_session.publish.return_value = mocker.Mock()
     publish_handle = client.publish(
         "test/topic",
-        b"test_payload",
+        payload,
         qos=qos,
         retain=True,
         properties=MQTTPublishProps(
@@ -98,9 +99,10 @@ def test_client_publish(qos: int | MQTTQoS, mocker: MockerFixture, mock_connecti
     )
     assert publish_handle == mock_session.publish.return_value
     expected_qos = MQTTQoS(qos) if not isinstance(qos, MQTTQoS) else qos
+    expected_payload = payload.encode("utf-8") if isinstance(payload, str) else bytes(payload)
     mock_session.publish.assert_called_once_with(
         "test/topic",
-        b"test_payload",
+        expected_payload,
         qos=expected_qos,
         retain=True,
         properties=MQTTPublishProps(