new mock interface in dask-awkward (#105)

dask-contrib · Oct 6, 2023 · fda6a1c · fda6a1c
1 parent a6d71b7
commit fda6a1c
Showing 1 changed file with 9 additions and 12 deletions.
diff --git a/src/dask_histogram/layers.py b/src/dask_histogram/layers.py
@@ -3,16 +3,13 @@
 
 class MockableDataFrameTreeReduction(DataFrameTreeReduction):
     def mock(self):
-        return (
-            MockableDataFrameTreeReduction(
-                name=self.name,
-                name_input=self.name_input,
-                npartitions_input=1,
-                concat_func=self.concat_func,
-                tree_node_func=self.tree_node_func,
-                finalize_func=self.finalize_func,
-                split_every=self.split_every,
-                tree_node_name=self.tree_node_name,
-            ),
-            None,
+        return MockableDataFrameTreeReduction(
+            name=self.name,
+            name_input=self.name_input,
+            npartitions_input=1,
+            concat_func=self.concat_func,
+            tree_node_func=self.tree_node_func,
+            finalize_func=self.finalize_func,
+            split_every=self.split_every,
+            tree_node_name=self.tree_node_name,
         )