Skip to content

Commit

Permalink
perf(common): reduce the average recusrion depth in `_flatten_collect…
Browse files Browse the repository at this point in the history
…ions` (#8709)
  • Loading branch information
kszucs authored Mar 21, 2024
1 parent 6639163 commit 3d52904
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 14 deletions.
15 changes: 7 additions & 8 deletions ibis/common/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,19 +56,18 @@ def _flatten_collections(node: Any) -> Iterator[N]:
>>> c = MyNode(2, "c", (a, b))
>>> d = MyNode(1, "d", (c,))
>>>
>>> assert list(_flatten_collections(a)) == [a]
>>> assert list(_flatten_collections((c,))) == [c]
>>> assert list(_flatten_collections([a, b, (c, a)])) == [a, b, c, a]
>>> assert list(_flatten_collections([{"b": b, "a": a}])) == [b, a]
"""
if isinstance(node, Node):
yield node
elif isinstance(node, (tuple, list)):
for item in node:
for item in node:
if isinstance(item, Node):
yield item
elif isinstance(item, (tuple, list)):
yield from _flatten_collections(item)
elif isinstance(node, (dict, frozendict)):
for value in node.values():
yield from _flatten_collections(value)
elif isinstance(item, (dict, frozendict)):
yield from _flatten_collections(item.values())


def _recursive_lookup(obj: Any, dct: dict) -> Any:
Expand Down
7 changes: 1 addition & 6 deletions ibis/common/tests/test_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -288,12 +288,7 @@ def test_flatten_collections():
assert list(result) == [A, B, C, D, E]

result = _flatten_collections(
{
"a": 0.0,
"b": A,
"c": (MyMapping(d=B, e=3), frozendict(f=C)),
"d": [5, "6", {"e": (D, 8.9)}],
}
[0.0, A, (MyMapping(d=B, e=3), frozendict(f=C)), [5, "6", {"e": (D, 8.9)}]]
)
assert list(result) == [A, C, D]

Expand Down

0 comments on commit 3d52904

Please sign in to comment.