use a context manager to automatically close the socket of UDPClient. Closes #193

Author: attwad

pythonosc/test/test_udp_client.py

@@ -48,5 +48,21 @@ def test_send_message_calls_add_arg_multiple_times_with_list(self):
         self.assertEqual(self.builder.add_arg.call_count, 3)
 
 
+class TestUdpClientClose(unittest.TestCase):
+    @mock.patch("socket.socket")
+    def test_close(self, mock_socket_ctor):
+        mock_socket = mock_socket_ctor.return_value
+        client = udp_client.UDPClient("::1", 31337)
+        client.close()
+        self.assertTrue(mock_socket.close.called)
+
+    @mock.patch("socket.socket")
+    def test_context_manager(self, mock_socket_ctor):
+        mock_socket = mock_socket_ctor.return_value
+        with udp_client.UDPClient("::1", 31337) as client:
+            self.assertIsInstance(client, udp_client.UDPClient)
+        self.assertTrue(mock_socket.close.called)
+
+
 if __name__ == "__main__":
     unittest.main()

pythonosc/udp_client.py

@@ -55,6 +55,16 @@ def __init__(
         self._address = address
         self._port = port
 
+    def __enter__(self) -> "UDPClient":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
+    def close(self) -> None:
+        """Close the socket"""
+        self._sock.close()
+
     def send(self, content: Union[OscMessage, OscBundle]) -> None:
         """Sends an :class:`OscMessage` or :class:`OscBundle` via UDP