mock interface correction (#89)

dask-contrib · Jun 13, 2023 · 5e7c033 · 5e7c033
1 parent f637e82
commit 5e7c033
Showing 1 changed file with 12 additions and 9 deletions.
diff --git a/src/dask_histogram/layers.py b/src/dask_histogram/layers.py
@@ -3,13 +3,16 @@
 
 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,
+        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,
         )