add tests

Author: siliataider

bindings/pyroot/pythonizations/test/rdf_define_pyz.py

@@ -113,6 +113,107 @@ def test_std_function(self):
         for x,y in zip(rdf2.Take['ULong64_t']("rdfentry_"), rdf2.Take['ULong64_t']("x")):
            self.assertEqual(x*x, y)
 
-    
+    def test_cpp_free_function(self):
+        """
+        Test that a C++ free function can be passed as a callable argument of a
+        Define operation.
+        """
+
+        test_cases = [
+            # Free function with arguments
+            {
+                "name": "input_ULong64_t",
+                "decl": """
+                    ULong64_t my_free_function(ULong64_t l) { return l; }
+                """,
+                "input": True,
+                "coltype": "ULong64_t",
+                "callable": lambda: ROOT.my_free_function,
+                "extract_fn": lambda x: x,
+            },
+            # Free function with user defined struct
+            {
+                "name": "input_user_defined_struct",
+                "decl": """
+                    struct MyStruct {
+                        ULong64_t value;
+                    };
+
+                    MyStruct my_free_function_struct(ULong64_t x) {
+                        MyStruct s;
+                        s.value = x;
+                        return s;
+                    }
+                """,
+                "input": True,
+                "coltype": "MyStruct",
+                "callable": lambda: ROOT.my_free_function_struct,
+                "extract_fn": lambda s: s.value,
+            },
+            # Free function with no arguments
+            {
+                "name": "no_input",
+                "decl": """
+                    ULong64_t my_free_function_none() { return 0; }
+                """,
+                "input": False,
+                "coltype": "ULong64_t",
+                "callable": lambda: ROOT.my_free_function_none,
+                "extract_fn": lambda x: x,
+            },
+        ]
+
+        for case in test_cases:
+            with self.subTest(case=case["name"]):
+                ROOT.gInterpreter.Declare(case["decl"])
+
+                rdf = ROOT.RDataFrame(5)
+                rdf = rdf.Define("new_col", case["callable"](), ["rdfentry_"]) if case["input"] else rdf.Define("new_col", case["callable"]())
+
+                inputs = rdf.Take["ULong64_t"]("rdfentry_") if case["input"] else [0] * 5
+                outputs = rdf.Take[case["coltype"]]("new_col")
+
+                for x, y in zip(inputs, outputs):
+                    self.assertEqual(case["extract_fn"](y), x)
+
+    def test_cpp_free_function_overloead(self):
+        """
+        Test that an overload of a C++ free function can be passed as a callable argument of a
+        Define operation with overloads.
+        """
+
+        ROOT.gInterpreter.Declare("""
+            ULong64_t my_free_function_overload(ULong64_t l) { return l; }
+            ULong64_t my_free_function_overload(ULong64_t l, ULong64_t m) { return l * m; }
+        """)
+
+        rdf = ROOT.RDataFrame(5)
+        rdf = rdf.Define("new_col", ROOT.my_free_function_overload, ["rdfentry_"])
+
+        for x, y in zip(rdf.Take["ULong64_t"]("rdfentry_"), rdf.Take["ULong64_t"]("new_col")):
+            self.assertEqual(x, y)
+
+        rdf = rdf.Define("new_col_overload", ROOT.my_free_function_overload, ["rdfentry_", "rdfentry_"])
+        for x, y in zip(rdf.Take["ULong64_t"]("rdfentry_"), rdf.Take["ULong64_t"]("new_col_overload")):
+            self.assertEqual(x * x, y)
+
+    def test_cpp_free_function_template(self):
+        """
+        Test that a templated C++ free function can be passed as a callable argument of a
+        Define operation.
+        """
+
+        ROOT.gInterpreter.Declare("""
+            template <typename T>
+            T my_free_function_template(T l) { return l; }
+        """)
+
+        rdf = ROOT.RDataFrame(5)
+        rdf = rdf.Define("new_col", ROOT.my_free_function_template["ULong64_t"], ["rdfentry_"])
+
+        for x, y in zip(rdf.Take["ULong64_t"]("rdfentry_"), rdf.Take["ULong64_t"]("new_col")):
+            self.assertEqual(x, y)
+
+
 if __name__ == '__main__':
     unittest.main()