From 714cd82d1b8a9d86739f9cb267405ac31fb26dc9 Mon Sep 17 00:00:00 2001 From: Martin Wendt Date: Sun, 22 Sep 2024 17:32:58 +0200 Subject: [PATCH] Optimite `visit()` performance --- nutree/common.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/nutree/common.py b/nutree/common.py index 9d57ed4..4face0c 100644 --- a/nutree/common.py +++ b/nutree/common.py @@ -297,7 +297,10 @@ def call_traversal_cb(fn: Callable, node: Node, memo: Any) -> False | None: """ try: res = fn(node, memo) - if res is SkipBranch or isinstance(res, SkipBranch): + + if res is None: + return None + elif res is SkipBranch or isinstance(res, SkipBranch): return False elif res is StopTraversal or isinstance(res, StopTraversal): raise res @@ -306,10 +309,10 @@ def call_traversal_cb(fn: Callable, node: Node, memo: Any) -> False | None: elif res is StopIteration or isinstance(res, StopIteration): # Converts wrong syntax in exception handler... raise res - elif res is not None: + else: raise ValueError( "callback should not return values except for " - f"False, SkipBranch, or StopTraversal: {res!r}." + f"None, False, SkipBranch, or StopTraversal: {res!r}." ) except SkipBranch: return False