fix test_mask

ASEM000 · Apr 2, 2024 · 249ae02 · 249ae02
1 parent 13d26b8
commit 249ae02
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/sepes/_src/tree_mask.py b/sepes/_src/tree_mask.py
@@ -245,7 +245,12 @@ def tree_mask(
         >>> sp.tree_unmask(masked_tree)
         [MyInt(value=1), MyInt(value=2), {'a': MyInt(value=3)}]
     """
-    return _tree_mask_map(tree, cond=cond, func=tree_mask.dispatcher, is_leaf=is_leaf)
+    return _tree_mask_map(
+        tree,
+        cond=cond,
+        func=tree_mask.dispatcher,
+        is_leaf=is_leaf,
+    )
 
 
 tree_mask.dispatcher = ft.singledispatch(_MaskedHashable)
@@ -314,10 +319,7 @@ def tree_unmask(tree: T, cond: Callable[[Any], bool] = lambda _: True):
         [MyInt(value=1), MyInt(value=2), {'a': MyInt(value=3)}]
     """
     return _tree_mask_map(
-        tree,
-        cond=cond,
-        func=tree_unmask.dispatcher,
-        is_leaf=is_masked,
+        tree, cond=cond, func=tree_unmask.dispatcher, is_leaf=is_masked
     )
 
 

diff --git a/tests/test_mask.py b/tests/test_mask.py
@@ -24,6 +24,7 @@
 from sepes._src.tree_base import TreeClass
 from sepes._src.tree_mask import is_masked, tree_mask, tree_unmask
 from sepes._src.tree_util import is_tree_equal, leafwise, tree_hash
+import dataclasses as dc
 
 test_arraylib = os.environ.get("SEPES_TEST_ARRAYLIB", "numpy")
 freeze = ft.partial(tree_mask, cond=lambda _: True)