Move maybe_add_argument to functional.py

Author: mdbenito

src/pydvl/utils/functional.py

@@ -4,7 +4,7 @@
 from functools import partial
 from typing import Callable, Set, Tuple, Union
 
-__all__ = ["get_free_args_fn", "fn_accept_additional_argument"]
+__all__ = ["maybe_add_argument"]
 
 
 def fn_accept_additional_argument(*args, fn: Callable, arg: str, **kwargs):
@@ -95,3 +95,25 @@ def _rec_unroll_partial_function_args(g: Union[Callable, partial]) -> Callable:
     wrapped_fn = _rec_unroll_partial_function_args(fun)
     sig = inspect.signature(wrapped_fn)
     return args_set_by_partial | set(sig.parameters.keys())
+
+
+def maybe_add_argument(fun: Callable, new_arg: str) -> Callable:
+    """Wraps a function to accept the given keyword parameter if it doesn't
+    already.
+
+    If `fun` already takes a keyword parameter of name `new_arg`, then it is
+    returned as is. Otherwise, a wrapper is returned which merely ignores the
+    argument.
+
+    Args:
+        fun: The function to wrap
+        new_arg: The name of the argument that the new function will accept
+            (and ignore).
+
+    Returns:
+        A new function accepting one more keyword argument.
+    """
+    if new_arg in free_arguments(fun):
+        return fun
+
+    return functools.partial(_accept_additional_argument, fun=fun, arg=new_arg)

src/pydvl/utils/parallel/map_reduce.py

@@ -20,8 +20,8 @@
     ReduceFunction,
     Seed,
     ensure_seed_sequence,
-    maybe_add_argument,
 )
+from ..functional import maybe_add_argument
 from .backend import init_parallel_backend
 
 __all__ = ["MapReduceJob"]

src/pydvl/utils/types.py

@@ -3,15 +3,12 @@
 """
 from __future__ import annotations
 
-import functools
 from abc import ABCMeta
-from typing import Any, Callable, Optional, Protocol, TypeVar, Union, cast
+from typing import Any, Optional, Protocol, TypeVar, Union, cast
 
 from numpy.random import Generator, SeedSequence
 from numpy.typing import NDArray
 
-from pydvl.utils.functional import fn_accept_additional_argument, get_free_args_fn
-
 __all__ = ["SupervisedModel", "MapFunction", "ReduceFunction", "NoPublicConstructor"]
 
 R = TypeVar("R", covariant=True)
@@ -45,28 +42,6 @@ def score(self, x: NDArray, y: NDArray) -> float:
         pass
 
 
-def maybe_add_argument(fun: Callable, new_arg: str) -> Callable:
-    """Wraps a function to accept the given keyword parameter if it doesn't
-    already.
-
-    If `fun` already takes a keyword parameter of name `new_arg`, then it is
-    returned as is. Otherwise, a wrapper is returned which merely ignores the
-    argument.
-
-    Args:
-        fun: The function to wrap
-        new_arg: The name of the argument that the new function will accept
-            (and ignore).
-
-    Returns:
-        A new function accepting one more keyword argument.
-    """
-    if new_arg in get_free_args_fn(fun):
-        return fun
-
-    return functools.partial(fn_accept_additional_argument, fn=fun, arg=new_arg)
-
-
 class NoPublicConstructor(ABCMeta):
     """Metaclass that ensures a private constructor