Better way of "find_all" to a certain depth

[big-c-note]

Thanks for the library, it's been awesome working with it. Wanted to pick someone's brain on a better way to do this

def bfs_limited(
    root: exp.Expression,
    max_depth: int = 1,
    match_types: t.Tuple[t.Type[exp.Expression]] = (),
) -> t.Iterator[exp.Expression]:
    # TODO: push this as helper to sqlglot or see if it can be done with walk
    # and depth.
    queue = deque([(root, 0)])  # node and its depth

    while queue:
        node, depth = queue.popleft()

        if depth > max_depth:
            continue

        if depth > 0:  # skip root node itself
            if not match_types or isinstance(node, match_types):
                yield node

        # only add children if we haven't reached max depth yet
        if depth < max_depth:
            queue.extend((child, depth + 1) for child in node.iter_expressions())

I had a need to find all up to a certain depth. A good example is when I need to get all columns that are selected in an exp.Select. So I have a node of type exp.CTE or exp.Select and I need to get the next exp.Columns. It's a BFS traversal, essentially I could think of two ways to do this.

1. Add a depth argument when using bfs traversal to find_all, and return all found of a certain type up to that depth.
2. Add a helper that finds the next exp.Column or a variable type, then return all nodes of that type at that depth.

Curious if you have interest in either of these, and I can submit a PR. Otherwise, is there a better way to do this with the library by passing something to prune in walk. Maybe that is the most elegant approach. Thank you in advance

[georgesittas]

Hey @big-c-note, based on your CTE & Select example, it appears that you may find the scope module useful. There are some examples in the AST primer you can look into.

If that doesn't fully cover your use case, I'd consider modifying the existing DFS logic to incorporate a depth counter. I don't think we're interested in merging this right now, as there's no concrete use case to justify the additional complexity.

Feel free to reach out in our Slack, which is more active compared to GitHub's "Discussions".

[big-c-note]

Got it thank you for the response