Overload unmap to accept handlers object or function

Author: DrLuke

pythonosc/dispatcher.py

@@ -4,6 +4,8 @@
 import re
 import time
 from pythonosc import osc_packet
+from typing import overload
+from types import FunctionType
 
 class Handler(object):
   def __init__(self, _callback, _args, _needs_reply_address=False):
@@ -58,7 +60,18 @@ def map(self, address, handler, *args, needs_reply_address=False):
     self._map[address].append(handlerobj)
     return handlerobj
 
-  def unmap(self, address, handler, *args, needs_reply_address=False):
+  @overload
+  def unmap(self, address: str, handler: Handler):
+    """Remove an already mapped handler from an address
+
+        Args:
+          - address: An explicit endpoint.
+          - handler: A Handler object as returned from map().
+    """
+    pass
+
+  @overload
+  def unmap(self, address: str, handler: FunctionType, *args, needs_reply_address: bool=False):
     """Remove an already mapped handler from an address
 
     Args:
@@ -68,8 +81,15 @@ def unmap(self, address, handler, *args, needs_reply_address=False):
       - args: Any additional arguments that will be always passed to the
               handlers after the osc messages arguments if any.
       - needs_reply_address: True if the handler function needs the
-              originating client address passed (as the first argument)."""
-    self._map[address].remove(Handler(handler, list(args), needs_reply_address))
+              originating client address passed (as the first argument).
+    """
+    pass
+
+  def unmap(self, address, handler, *args, needs_reply_address=False):
+    if isinstance(handler, Handler):
+      self._map[address].remove(handler)
+    else:
+      self._map[address].remove(Handler(handler, list(args), needs_reply_address))
 
   def handlers_for_address(self, address_pattern):
     """yields Handler namedtuples matching the given OSC pattern."""

pythonosc/test/test_dispatcher.py

@@ -125,10 +125,17 @@ def test_unmap(self):
     def dummyhandler():
         pass
 
-    self.dispatcher.map("/map/me", dummyhandler)
+    # Test with handler returned by map
+    returnedhandler = self.dispatcher.map("/map/me", dummyhandler)
     self.sortAndAssertSequenceEqual([Handler(dummyhandler, [])], self.dispatcher.handlers_for_address("/map/me"))
-    self.dispatcher.unmap("/map/me", dummyhandler)
+    self.dispatcher.unmap("/map/me", returnedhandler)
     self.sortAndAssertSequenceEqual([], self.dispatcher.handlers_for_address("/map/me"))
 
+    # Test with reconstructing handler
+    self.dispatcher.map("/map/me/too", dummyhandler)
+    self.sortAndAssertSequenceEqual([Handler(dummyhandler, [])], self.dispatcher.handlers_for_address("/map/me/too"))
+    self.dispatcher.unmap("/map/me/too", dummyhandler)
+    self.sortAndAssertSequenceEqual([], self.dispatcher.handlers_for_address("/map/me/too"))
+
 if __name__ == "__main__":
   unittest.main()