add leaf_class to decorators

Author: majosm

arraycontext/container/traversal.py

@@ -268,12 +268,31 @@ def rec_map_array_container(
     return _map_array_container_impl(f, ary, leaf_cls=leaf_class, recursive=True)
 
 
+def _mapped_over_array_containers_factory(
+        leaf_class: Optional[type] = None) -> Callable[[Callable[
+            [Any], Any]], Callable[[ArrayOrContainerT], ArrayOrContainerT]]:
+    """Decorator factory around :func:`rec_map_array_container`."""
+    def decorator(f: Callable[[Any], Any]) -> Callable[
+            [ArrayOrContainerT], ArrayOrContainerT]:
+        wrapper = partial(rec_map_array_container, f, leaf_class=leaf_class)
+        update_wrapper(wrapper, f)
+        return wrapper
+    return decorator
+
+
 def mapped_over_array_containers(
-        f: Callable[[Any], Any]) -> Callable[[ArrayOrContainerT], ArrayOrContainerT]:
+        f: Optional[Callable[[Any], Any]] = None,
+        leaf_class: Optional[type] = None) -> Union[
+            Callable[[ArrayOrContainerT], ArrayOrContainerT],
+            Callable[
+                Callable[[Any], Any],
+                Callable[[ArrayOrContainerT], ArrayOrContainerT]]]:
     """Decorator around :func:`rec_map_array_container`."""
-    wrapper = partial(rec_map_array_container, f)
-    update_wrapper(wrapper, f)
-    return wrapper
+    decorator = _mapped_over_array_containers_factory(leaf_class=leaf_class)
+    if f is not None:
+        return decorator(f)
+    else:
+        return decorator
 
 
 def rec_multimap_array_container(
@@ -291,16 +310,32 @@ def rec_multimap_array_container(
         f, *args, leaf_cls=leaf_class, recursive=True)
 
 
+def _multimapped_over_array_containers_factory(
+        leaf_class: Optional[type] = None) -> Callable[
+            [Callable[..., Any]], Callable[..., Any]]:
+    """Decorator factory around :func:`rec_multimap_array_container`."""
+    def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+        # can't use functools.partial, because its result is insufficiently
+        # function-y to be used as a method definition.
+        def wrapper(*args: Any) -> Any:
+            return rec_multimap_array_container(f, *args, leaf_class=leaf_class)
+        update_wrapper(wrapper, f)
+        return wrapper
+    return decorator
+
+
 def multimapped_over_array_containers(
-        f: Callable[..., Any]) -> Callable[..., Any]:
+        f: Optional[Callable[..., Any]] = None,
+        leaf_class: Optional[type] = None) -> Union[
+            Callable[..., Any],
+            Callable[[Callable[..., Any]], Callable[..., Any]]]:
     """Decorator around :func:`rec_multimap_array_container`."""
-    # can't use functools.partial, because its result is insufficiently
-    # function-y to be used as a method definition.
-    def wrapper(*args: Any) -> Any:
-        return rec_multimap_array_container(f, *args)
+    decorator = _multimapped_over_array_containers_factory(leaf_class=leaf_class)
+    if f is not None:
+        return decorator(f)
+    else:
+        return decorator
 
-    update_wrapper(wrapper, f)
-    return wrapper
 
 # }}}