Add varargs support to multiple dispatch

Technologicat · Technologicat · commit cbc4641ef460 · 2020-03-20T14:50:22.000+02:00
diff --git a/unpythonic/dispatch.py b/unpythonic/dispatch.py
@@ -21,6 +21,7 @@
 __all__ = ["generic", "specific"]
 
 from functools import wraps
+from itertools import chain
 import inspect
 import typing
 
@@ -83,9 +84,15 @@ def generic(f):
     def multidispatch(*args, **kwargs):
         # signature comes from typing.get_type_hints.
         def match_argument_types(signature):
-            # TODO: handle *args (bindings["vararg"], bindings["vararg_name"])
             # TODO: handle **kwargs (bindings["kwarg"], bindings["kwarg_name"])
-            for parameter, value in bindings["args"].items():
+            args_items = bindings["args"].items()
+            if bindings["vararg_name"]:
+                vararg_item = (bindings["vararg_name"], bindings["vararg"])  # *args
+                all_items = tuple(chain(args_items, (vararg_item,)))
+            else:
+                all_items = args_items
+
+            for parameter, value in all_items:
                 assert parameter in signature  # resolve_bindings should already TypeError when not.
                 expected_type = signature[parameter]
                 if not match_value_to_typespec(value, expected_type):
diff --git a/unpythonic/test/test_dispatch.py b/unpythonic/test/test_dispatch.py
@@ -35,6 +35,19 @@ def example(start: int, step: int, stop: int):
 def _example_impl(start, step, stop):
     return start, step, stop
 
+# varargs are supported via `typing.Tuple`
+@generic
+def gargle(): ...
+@gargle.register
+def gargle(*args: typing.Tuple[int, ...]):  # any number of ints
+    return "int"
+@gargle.register
+def gargle(*args: typing.Tuple[float, ...]):  # any number of floats
+    return "float"
+@gargle.register
+def gargle(*args: typing.Tuple[int, float, str]):  # three args, matching the given types
+    return "int, float, str"
+
 # One-method pony.
 @specific
 def blubnify(x: int, y: float):
@@ -62,6 +75,11 @@ def test():
     assert example(1, 1, 5) == (1, 1, 5)
     assert example(1, 2, 5) == (1, 2, 5)
 
+    assert gargle(1, 2, 3, 4, 5) == "int"
+    assert gargle(2.71828, 3.14159) == "float"
+    assert gargle(42, 6.022e23, "hello") == "int, float, str"
+    assert gargle(1, 2, 3) == "int"  # as many as in the [int, float, str] case
+
     assert blubnify(2, 21.0) == 42
     try:
         blubnify(2, 3)
