The command line tool is a thin wrapper around the `TreeContext` class, whose purpose is show you a set of "lines of interest" in the context of the entire AST. This all exists because my other project aider [0] uses TreeContext to display a repository map [1] so that GPT-4 can understand how the most important classes, methods, functions, etc fit into the entire code base of a git repository.
But it was easy to make a CLI interface to grep lines of interest and display them with TreeContext, and it turned out to be quite useful.
The TreeContext class is line-oriented, and is mainly interested in tracking language constructs whose scope spans multiple lines. Typically these are things like classes, methods, functions, loops, if/else constructs, etc. Given a line of interest, we look at all the multi-line scopes which contain it. For each such multi-line scope, we want to display some "header" lines to provide context.
In this example, the match for "two" is contained in the multi-line scopes of a method and a class. So we print their headers.
$ grep-ast two example.py
⋮...
│class MyClass:
│ "MyClass is great"
⋮...
│ def print2(self):
▶ print("two")
⋮...
For any multi-line scope, we look at all the other AST scopes which start on the same line. We take the smallest such co-occurring scope, and declare the header to be the lines that it spans. For a simple method like `def print2(self):`, that's all that gets picked up.
But a complex method like `print1()` below picks up all the lines which are part of its full function signature:
$ grep-ast one example.py
⋮...
│class MyClass:
│ "MyClass is great"
⋮...
│ def print1(
│ self,
│ prefix,
│ suffix,
│ ):
⋮...
▶ print(f"{prefix} one {suffix}")
⋮...