From 0ac8b394576244fe6675df0c018565c1beb7b384 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 22 Nov 2021 13:09:53 -0800
Subject: [PATCH] Add full support for new reduction functions to match
 docstrings and signatures.

---
 python/cudf/cudf/core/reductions.py | 61 +++++++++++++++++++++++++----
 1 file changed, 54 insertions(+), 7 deletions(-)

diff --git a/python/cudf/cudf/core/reductions.py b/python/cudf/cudf/core/reductions.py
index e554dd34e35..c7c8b10118e 100644
--- a/python/cudf/cudf/core/reductions.py
+++ b/python/cudf/cudf/core/reductions.py
@@ -6,13 +6,20 @@
 # for reductions, scans, and binops (and perhaps others).
 
 
-# TODO: The docstring for a class's reductions should be taken from formatting
-# the reduce method since it needs to support all the same things.
+# TODO: Figure out how to support examples in the docstrings.
 
 
 # TODO: Consider using a pyi file to trick mypy into seeing the monkey-patched
 # methods.
 
+# TODO: Add type annotations for the class variables.
+
+# TODO: Extract signature from _reduce
+
+# TODO: Take advantage of the new approach to also format the cls into the doc.
+
+import inspect
+
 
 class Reducible:
     """Mixin encapsulating for reduction operations.
@@ -70,12 +77,52 @@ def _reduce(self, op: str, *args, **kwargs):
 
     @classmethod
     def _add_reduction(cls, reduction):
-        def op(self, *args, **kwargs):
-            return self._reduce(reduction, *args, **kwargs)
+        # This function creates reduction operations on-the-fly and assigns
+        # them to the class.
+
+        # Generate a signature without the `op` parameter.
+        signature = inspect.signature(cls._reduce)
+        new_params = signature.parameters.copy()
+        new_params.pop("op")
+        signature = signature.replace(parameters=new_params.values())
+
+        # Generate the list of arguments forwarded to _reduce.
+        arglist = ", ".join(
+            [
+                f"{key}={key}"
+                for key in signature.parameters
+                if key not in ("self", "args", "kwargs")
+            ]
+        )
+        if arglist:
+            arglist += ", *args, **kwargs"
+        else:
+            arglist = "*args, **kwargs"
+
+        # The default docstring is that of the _reduce method. Additional
+        # formatting arguments may be provided in a class-level dictionary
+        # of the form _REDUCTION_DOCSTRINGS
+        docstring = cls._reduce.__doc__.format(
+            cls=cls.__name__,
+            op=reduction,
+            **getattr(cls, "_REDUCTION_DOCSTRINGS", {}).get(reduction, {}),
+        )
 
-        # The default docstring is that
-        op.__doc__ = cls._reduce.__doc__.format(op=reduction)
-        setattr(cls, reduction, op)
+        # Create the desired function.
+        namespace = {}
+        out = """
+def {reduction}{signature}:
+    \"\"\"{docstring}
+    \"\"\"
+    return self._reduce(op="{reduction}", {arglist})
+        """.format(
+            reduction=reduction,
+            signature=str(signature),
+            arglist=arglist,
+            docstring=docstring,
+        )
+        exec(out, namespace)
+        setattr(cls, reduction, namespace[reduction])
 
     @classmethod
     def __init_subclass__(cls):