python header ag -- search for expression and get nearest headers

header-ag.py

"""
Usage:

    python my-ag.py <regex> <...paths>

Requirements/Features of my-ag.py:
* Takes regex pattern and file/directory paths as command-line arguments
* Recursively searches through Python files in given paths
* When finding a regex match, shows:
  - All parent unindented lines ("headers")
  - Proper indentation hierarchy
  - "..." between indent levels
  - Line numbers
  - Colored output (filepath, matches, line numbers)
* Handles both individual files and directories as input
* Preserves original indentation structure in output
* Short & simple
* No argparse
"""
import sys
import re
import glob
import os
from termcolor import colored

def find_python_files(paths):
    files = []
    for path in paths:
        if os.path.isfile(path) and path.endswith('.py'):
            files.append(path)
        elif os.path.isdir(path):
            files.extend(glob.glob(f"{path}/**/*.py", recursive=True))
    return files

def process_file(filepath, pattern, context=0):
    with open(filepath) as f:
        lines = f.readlines()

    headers = []  # Stack of (indent_level, line, line_num) tuples
    matches = []  # List of (line_num, headers, line) tuples for matches

    for i, line in enumerate(lines):
        stripped = line.strip()
        if not stripped:  # Skip empty lines
            continue

        indent = len(line) - len(line.lstrip())

        # Update headers stack if this is a new header
        if indent == 0:
            headers = [(0, line.rstrip(), i + 1)]
        else:
            # Remove any headers that are at same or deeper indent
            while headers and headers[-1][0] >= indent:
                headers.pop()
            headers.append((indent, line.rstrip(), i + 1))

        # Check for regex match
        if re.search(pattern, line):
            matches.append((i, headers[:], line.rstrip()))

    if matches:
        print(f"\n{colored(filepath, 'cyan')}:")

        for match_idx, (match_line_num, match_headers, match_line) in enumerate(matches):
            # Print separator between match groups
            if match_idx > 0:
                print("--")

            # Print context before
            start_line = max(0, match_line_num - context)
            for i in range(start_line, match_line_num):
                if lines[i].strip():
                    line_prefix = colored(f"{i+1:4d}", "yellow") + " "
                    print(line_prefix + lines[i].rstrip())

            # Print headers and match
            last_num = -1
            for h_indent, h_line, h_num in match_headers[:-1]:
                if last_num >= 0 and h_num - last_num > 1:
                    print("     " + " " * h_indent + colored("...", "grey"))
                line_prefix = colored(f"{h_num:4d}", "yellow") + " "
                if h_indent == 0:
                    print('     ***')
                    h_line = colored(h_line, "green")
                print(line_prefix + h_line)
                last_num = h_num

            matched_num = match_headers[-1][2]
            matched_line = match_headers[-1][1]
            if last_num >= 0 and matched_num - last_num > 1:
                numspaces = len(matched_line) - len(matched_line.lstrip())
                print("     " + " " * numspaces + colored("...", "grey"))

            line_prefix = colored(f"{matched_num:4d}", "yellow") + " "
            # Highlight the matched pattern
            matched_line = re.sub(pattern, lambda m: colored(m.group(), 'red', attrs=['bold']), matched_line)
            print(line_prefix + matched_line)

            # Print context after
            end_line = min(len(lines), match_line_num + context + 1)
            for i in range(match_line_num + 1, end_line):
                if lines[i].strip():
                    line_prefix = colored(f"{i+1:4d}", "yellow") + " "
                    print(line_prefix + lines[i].rstrip())

def main():
    if len(sys.argv) < 3:
        print("Usage: python my-ag.py [--context=num] <regex> <...paths>")
        sys.exit(1)

    context = 0
    args = sys.argv[1:]

    # Handle --context argument
    for i, arg in enumerate(list(args)):
        print(f"{i=} {arg=}")
        if arg.startswith('--context='):
            context = int(arg.split('=')[1])
            print(f"{context=}")
            args.pop(i)

    if len(args) < 2:
        print("Usage: python my-ag.py [--context=num] <regex> <...paths>")
        sys.exit(1)

    pattern = args[0]
    paths = args[1:]

    for file in find_python_files(paths):
        process_file(file, pattern, context)

if __name__ == "__main__":
    main()