[pyroot] Refactor RDF Filter/Define callable dispatch [ROOT-patch]

Author: aaronj0

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rdf_pyz.py

@@ -231,139 +231,30 @@ def jit_function(self, func, cols_list, extra_args):
         return self.func_call
 
 
-def remove_fn_name_from_signature(signature):
-    """
-    Removes the function name from a signature string.
-    The signature is expected to be in the form of "return_type function_name(type param1, type param2, ...)".
-    """
-    if "(" not in signature or ")" not in signature:
-        raise ValueError(f"Invalid signature format: {signature}")
-
-    return signature[: signature.index(" ")] + signature[signature.index("(") :]
-
-
-def get_cpp_overload_from_templ_proxy(func, types=None):
-    """
-    Gets the C++ overload from a cppyy TemplateProxy using the input
-    and template arguments types.
-    """
-    signature = ", ".join(types) if types else ""
-    template_args = func.__template_args__[1:-1] if func.__template_args__ else ""
-    return func.__overload__(signature, template_args)
-
-
 def get_column_types(rdf, cols):
     return [rdf.GetColumnType(col) for col in cols]
 
 
-def get_overload_based_on_args(overload_types, types):
-    """
-    Gets the C++ overloads for a function based on the provided signatures and types.
-    """
-    if not isinstance(overload_types, dict):
-        raise TypeError("Overload types must be a dictionary.")
-
-    if not isinstance(types, list) or not all(isinstance(t, str) for t in types):
-        raise TypeError("Types must be a list of strings.")
-
-    if len(overload_types) == 1:
-        # If there is only one signature, return it directly
-        return next(iter(overload_types))
-
-    for full_sig, info in overload_types.items():
-        input_types = info.get("input_types", ())
-
-        if len(input_types) != len(types):
-            continue
-
-        match = True
-        for expected, actual in zip(input_types, types):
-            # Loose match: require actual to appear somewhere in expected (e.g. "int" in "const int&")
-            if actual not in expected:
-                match = False
-                break
-
-        if match:
-            return full_sig
-
-    raise ValueError(
-        f"No matching overload found for types: {types} in signatures: {overload_types.keys()}. Please check the function overloads."
-    )
-
-
-def _get_cpp_signature(func, rdf, cols):
+def _resolve_overload_based_on_cols(func, rdf, cols):
     """
     Gets the C++ signature of a cppyy callable.
     """
     import ROOT
 
-    if isinstance(func, ROOT._cppyy.types.Template):
-        func = get_cpp_overload_from_templ_proxy(func, get_column_types(rdf, cols))
-
-    if not isinstance(func, ROOT._cppyy.types.Function):
+    if not isinstance(func, (ROOT._cppyy.types.Function, ROOT._cppyy.types.Template)):
         raise TypeError(f"Expected a cppyy callable, got {type(func).__name__}")
 
-    overload_types = func.func_overloads_types
-    matched_overload = get_overload_based_on_args(overload_types, get_column_types(rdf, cols))
-    return remove_fn_name_from_signature(matched_overload)
+    if isinstance(func, ROOT._cppyy.types.Template) and func.__template_args__:
+        template_args = func.__template_args__[1:-1]
+    else:
+        template_args = ""
 
-
-def _to_std_function(func, rdf, cols):
-    """
-    Converts a cppyy callable to std::function.
-    """
-    import ROOT
-
-    if not isinstance(func, ROOT._cppyy.types.Function) and not isinstance(func, ROOT._cppyy.types.Template):
-        raise TypeError(f"Expected a cppyy callable, got {type(func).__name__}")
-
-    signature = _get_cpp_signature(func, rdf, cols)
-    return ROOT.std.function(signature)
-
-
-def _handle_cpp_callables(func, original_template, *args, rdf=None, cols=None):
-    """
-    Checks whether the callable `func` is a cppyy proxy of one of these:
-    1. C++ functor
-    2. std::function
-    3. C++ free function
-
-    Cases 1 and 2 above are supported by cppyy, so we can just invoke the original
-    cppyy TemplateProxy (Filter or Define) with the callable as argument.
-    For case 3, we need to convert the callable to a std::function
-    before invoking the original cppyy TemplateProxy.
-
-    Prior to the invocation of the original cppyy TemplateProxy, though, we
-    need to explicitly instantiate the template using the type of the `func`
-    callable. Implicit instantiation won't work since the original
-    implementation of Filter and Define is replaced by a pythonization.
-
-    Arguments:
-        func: callable to be checked
-        original_rdf_method: cppyy proxy for Filter or Define
-        args: arguments passed by the user to Filter or Define
-    Returns:
-        RDataFrame: RDataFrame node if `func` can be classified in one of the
-    three cases above, None otherwise
-    """
-    import ROOT
-
-    is_cpp_functor = lambda: isinstance(getattr(func, "__call__", None), ROOT._cppyy.types.Function)  # noqa E731
-
-    is_std_function = lambda: isinstance(getattr(func, "target_type", None), ROOT._cppyy.types.Function)  # noqa E731
-
-    # handle free functions
-    if callable(func) and not is_cpp_functor() and not is_std_function():
-        try:
-            func = _to_std_function(func, rdf, cols)
-        except TypeError as e:
-            if "Expected a cppyy callable" in str(e):
-                pass  # this function is not convertible to std::function, move on to the next check
-            else:
-                raise
-
-    if is_cpp_functor() or is_std_function():
-        return original_template[type(func)](*args)
+    col_types = get_column_types(rdf, cols)
+    signature = ", ".join(col_types) if cols else ""
+    if template_args:
+        return func.__overload__(signature, template_args)
+    else:
+        return func.__overload__(signature)
 
 
 class _WarnOnce:
@@ -460,9 +351,13 @@ def x_more_than_y(x):
                 f"Arguments should be ('list', 'str',) not ({type(args[0]).__name__, type(args[1]).__name__}."
             )
 
-    rdf_node = _handle_cpp_callables(func, rdf._OriginalFilter, func, *args, rdf=rdf, cols=col_list)
-    if rdf_node is not None:
-        return rdf_node
+    import ROOT
+
+    if isinstance(func, (ROOT._cppyy.types.Function, ROOT._cppyy.types.Template)):
+        return rdf._OriginalFilter(_resolve_overload_based_on_cols(func, rdf, col_list), *args)
+
+    if isinstance(func, ROOT._cppyy.types.Instance):
+        return rdf._OriginalFilter(func, *args)
 
     _WarnOnce.warn()
     jitter = FunctionJitter(rdf)
@@ -520,9 +415,13 @@ def x_scaled(x):
         raise TypeError(f"Define takes a column list as third arguments but {type(cols).__name__} was given.")
 
     func = callable_or_str
-    rdf_node = _handle_cpp_callables(func, rdf._OriginalDefine, col_name, func, cols, rdf=rdf, cols=cols)
-    if rdf_node is not None:
-        return rdf_node
+
+    import ROOT
+
+    if isinstance(func, (ROOT._cppyy.types.Function, ROOT._cppyy.types.Template)):
+        return rdf._OriginalDefine(col_name, _resolve_overload_based_on_cols(func, rdf, cols), cols)
+    if isinstance(func, ROOT._cppyy.types.Instance):
+        return rdf._OriginalDefine(col_name, func, cols)
 
     _WarnOnce.warn()
     jitter = FunctionJitter(rdf)