I think this will only create the range for top-level import blocks as it doesn't traverse into the stmt, so the following won't create the range:

def foo():
    import sys
    from math import prod

    def bar():
        pass

Can we use the single visitor with state handling instead? That would also means that we do one pass over the AST instead of multiple.

So I made the nested case work without moving state handling into the visitor. I added a TODO to consider moving it into the visitor, but I'm not totally convinced that it would be meaningfully faster. In particular, this is just doing a shallow scan across a sequence of statements (now all sequences of statements to handle the nested case), so it's not re-doing all of the AST visiting work. And if this were moved into the AST visitor, there would need to be more state tracking.

To be clear, I don't feel strong about this. My instincts here might be wrong!

We might not need this function? Because the visitor should visit these string expressions and it already has branch for strings and f-strings. We can expand them to cover bytes and t-strings as well.

That was my thought too. But we ideally need to distinguish between "normal" string literals and string literals that serve as documentation. In the latter case, we can attach the FoldingRangeKind::Comment kind. And that enables things like "collapse all comment blocks."

Oh, interesting. Is "collapse all comment blocks" an editor feature? This makes sense but I think I'd find it a bit confusing as to why did a function doc gets collapsed along with other comments but we can iterate on this based on user feedback.

Yeah! I'm just realizing that I forgot to post my demo video. Whoops.

Before:

After:

I think I'd be more surprised if "fold all comment blocks" didn't include docstrings personally. But yeah, happy to iterate in response to user feedback.

Can we use the single visitor with state handling instead? That would also means that we do one pass over the AST instead of multiple.