Improve multiple dispatch

Juha Jeronen · Juha Jeronen · commit e1ccfc8d7d08 · 2020-03-16T15:30:33.000+02:00
diff --git a/unpythonic/dispatch.py b/unpythonic/dispatch.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8; -*-
 """A multiple-dispatch decorator for Python.
 
-Somewhat like `functools.singledispatch`, but multiple dispatch.
+Somewhat like `functools.singledispatch`, but for multiple dispatch.
 
     https://docs.python.org/3/library/functools.html#functools.singledispatch
 
@@ -11,72 +11,88 @@
 
 **WARNING: PROVISIONAL FEATURE**
 
-This provisional feature is provided for technical preview purposes only.
+This provisional feature is a proof-of-concept provided for technical preview
+and teaching purposes only.
+
 Details may still change in a backwards-incompatible way, or the whole
 feature may still be removed. Do not depend on it in production!
 """
 
-from collections import defaultdict
+__all__ = ["generic", "specific"]
+
 from functools import wraps
 import inspect
 import typing
 
 from .arity import resolve_bindings
 
-registry = defaultdict(list)
-
 def generic(f):
     """Decorator. Make `f` a generic function (in the sense of CLOS or Julia).
 
-    Methods are attached by specifying type hints on parameters when defining
-    implementations. Then, at call time, types of **all** arguments are then
-    automatically used for choosing which method to call. Multiple parameters
-    may be used for dispatching.
+    **How to use**:
+
+    The `@generic`-decorated function definition itself *is just a declaration*
+    that defines the shape and parameter names of the formal parameter list.
+    The formal parameter list must remain the same across all methods of the
+    same generic function; only the types of the parameters may vary. That
+    function is otherwise a stub. It is never called.
+
+    Like in `functools.singledispatch`, methods (implementations for specific
+    combinations of argument types) are registered with the decorator
+    `@f.register`, where `f` is the function you decorated with `@generic`.
+
+    Each method must specify type hints **on all parameters**. Then, at call
+    time, types of **all** arguments are then automatically used for choosing
+    which method to call, i.e., multiple parameters are used for dispatching.
+
+    The first match wins, in most-recently-registered order. So specify the
+    implementation with the most generic types first, and then move on to the
+    more specific ones. The mnemonic is, "the function is generally defined
+    like this, except if the arguments match these particular types..."
 
     The point of this feature is to eliminate `if`/`elif`/`elif`... blocks
-    that switch by `isinstance` on arguments (and then raise `TypeError`
-    in the final `else`), by implementing the machinery once centrally.
+    that switch by `isinstance` on arguments, and then raise `TypeError`
+    in the final `else`, by implementing the machinery once centrally.
+
+    **Differences to tools in the standard library**:
 
-    Unlike `typing.overload`, the implementation belongs right there in the
-    type-specialized methods. Unlike `functools.singledispatch`, there is
-    no need to provide a fallback non-annotated implementation (and in fact
-    that is not supported).
+    Unlike `functools.singledispatch`, the `@generic` function itself is
+    unused, except as an interface declaration for the parameter list.
 
-    Upon ambiguity, the method that was most recently registered wins. So
-    specify the implementation with the most generic types first, and then
-    move on to the more specific ones. The mnemonic is, "the function is
-    generally defined like this, except if you get arguments that match
-    these particular types..."
+    Unlike `typing.overload`, the implementations are to be provided in the
+    method bodies.
 
     **CAUTION**:
 
-    Currently, the shape of the parameter list must agree between all methods
-    of the same generic function, and advanced features of `typing` are not
-    supported.
+    To declare a parameter of a method as dynamically typed, explicitly
+    annotate it as `typing.Any`; don't just omit the type annotation.
+    Explicit is better than implicit; this is a feature.
 
-    Use `typing.Any` to declare a dynamically typed parameter - don't just omit
-    the type annotation, that won't work.
+    Currently, advanced features of `typing` such as `Sequence[...]` are
+    not supported. This may or may not change in the future.
     """
-    fullname = "{}.{}".format(f.__module__, f.__qualname__)
-    signature = typing.get_type_hints(f)
-    registry[fullname].append((f, signature))
+    # Dispatcher - this will replace the original f.
     @wraps(f)
     def multidispatch(*args, **kwargs):
-        # TODO: This uses the definition of `f` this particular instance of the decorator
-        # TODO: was applied to. Maybe we should check, upon registration, that the shape
-        # TODO: of the parameter list matches for each method of the same generic function.
+        # We use the parameter list of the original `@generic`-decorated
+        # function to match the function call arguments to the formal
+        # parameters of the function definition.
         bindings = resolve_bindings(f, *args, **kwargs)
 
         def match(signature):
-            # TODO: handle *args (bindings["vararg"]) and **kwargs (bindings["kwarg"])
+            # TODO: handle *args (bindings["vararg"])
+            # TODO: handle **kwargs (bindings["kwarg"])
+            # TODO: handle advanced features such as Sequence[int], Optional[str], ...
             for parameter, value in bindings["args"].items():
                 p = signature[parameter]  # TODO: what if parameter is not there? TypeError?
                 if p is not typing.Any and not isinstance(value, p):
                     return False
             return True
 
-        methods = tuple(reversed(registry[fullname]))
-        for method, signature in methods:
+        # Dispatch.
+        def methods():
+            return reversed(multidispatch._registry)
+        for method, signature in methods():
             if match(signature):
                 return method(*args, **kwargs)
 
@@ -86,7 +102,7 @@ def match(signature):
         # TODO: but in the general case this is difficult. We can't just `type(x)`, since
         # TODO: the signature may specify something like `Sequence[int]`. Knowing a `list`
         # TODO: was passed doesn't help debug that it was `Sequence[str]` when a `Sequence[int]`
-        # TODO: was expected. The actual value at least contains the type information implicitly.
+        # TODO: was expected. The actual value at least implicitly contains the type information.
         #
         # TODO: Compute closest candidates, like Julia does? (see methods, MethodError)
         a = [str(a) for a in args]
@@ -97,31 +113,75 @@ def format_method(method):  # Taking a page from Julia and some artistic liberty
             filename = inspect.getsourcefile(obj)
             source, firstlineno = inspect.getsourcelines(obj)
             return "{} from {}:{}".format(signature, filename, firstlineno)
-        methods_str = ["  {}".format(format_method(x)) for x in methods]
+        methods_str = ["  {}".format(format_method(x)) for x in methods()]
         candidates = "\n".join(methods_str)
         msg = ("No method found matching {}({}{}{}).\n"
                "Candidate signatures (in order of match attempts):\n{}").format(f.__qualname__,
                                                                                 ", ".join(a),
                                                                                 sep, ", ".join(kw),
                                                                                 candidates)
         raise TypeError(msg)
+
+    # fullname = "{}.{}".format(f.__module__, f.__qualname__)
+    multidispatch._registry = []
+    def register(function):
+        """Decorator. Register a new method for this generic function.
+
+        The method must have type annotations for all of its parameters;
+        these are used for dispatching.
+        """
+        # TODO: fail-fast: verify the shape of the parameter list of `function`
+        # is compatible with the `@generic` function to which we are attaching
+        # `function` as a new method.
+        signature = typing.get_type_hints(function)
+        multidispatch._registry.append((function, signature))
+        return multidispatch  # Replace the function with the dispatcher for this generic function.
+    multidispatch.register = register  # publish the @f.register decorator
+
     return multidispatch
 
+
+def specific(f):
+    """Decorator. A one-method pony, which is kind of the opposite of `@generic`.
+
+    This restricts the allowed argument types to one combination only. This can
+    be used to eliminate `isinstance` boilerplate code in the function body, by
+    allowing the types (for dynamic, run-time checking) to be specified with a
+    very compact syntax.
+
+    Unlike in `@generic`, the function to be decorated simultaneously specifies
+    both the shape and parameter names of the formal parameter list, as well as
+    provides the body of the only method implementation.
+
+    A `@specific` function has no `.register` attribute; after it is created,
+    no more methods can be attached to it.
+    """
+    s = generic(f)
+    s.register(f)
+    del s.register  # remove the ability to attach more methods
+    return s
+
 # --------------------------------------------------------------------------------
 
 @generic
+def zorblify(x, y):  # could use the ellipsis `...` as the body, but this is a unit test.
+    assert False  # Stub, used only for interface declaration. Never called.
+@zorblify.register
 def zorblify(x: int, y: int):
     return 2 * x + y
-@generic
+@zorblify.register
 def zorblify(x: str, y: int):
-    # Because dispatching occurs on both arguments, this is not reached in tests.
-    assert False
-@generic
+    assert False  # Because dispatching occurs on both arguments, this is not reached in tests.
+@zorblify.register
 def zorblify(x: str, y: float):
     return "{} {}".format(x[::-1], y)
 
 # TODO: def zorblify(x: int, *args: typing.Sequence[str]):
 
+@specific
+def blubnify(x: int, y: float):
+    return x * y
+
 def test():
     assert zorblify(17, 8) == 42
     assert zorblify(17, y=8) == 42  # can also use named arguments
@@ -134,7 +194,16 @@ def test():
     except TypeError:
         pass
     else:
-        assert False  # should have noticed there's no method registered for zorblify(float, float)
+        assert False  # there's no zorblify(float, float)
+
+    assert blubnify(2, 21.0) == 42
+    try:
+        blubnify(2, 3)
+    except TypeError:
+        pass
+    else:
+        assert False  # blubnify only accepts (int, float)
+    assert not hasattr(blubnify, "register")
 
     print("All tests PASSED")
 
