Parse Python requirements files into AST JSON representation via tree-sitter-requirements

py-requirements-to-json

#!/usr/bin/env pip-run
"""
Parse Python requirements files into AST JSON representation
"""

# /// script
# dependencies = ['tree-sitter-requirements']
# ///

import json
import sys
from subprocess import check_output
from typing import Iterator, Tuple, Union

import tree_sitter_requirements as reqs
from tree_sitter import TreeCursor


# https://jqlang.github.io/jq/manual/#streaming
JSONValue = Union[None, bool, str, int, dict, list]
JSONPathItem = Union[int, str]
JSONPath = Tuple[JSONPathItem, ...]
JSONItem = Tuple[JSONPath, JSONValue]
JSONClose = Tuple[JSONPath]
JSONStream = Iterator[Union[JSONItem, JSONClose]]


def json_stream(cursor: TreeCursor, ctx: JSONPath = ()) -> JSONStream:
    i = 0
    yield ((*ctx, "type"), cursor.node.type)
    yield ((*ctx, "text"), cursor.node.text.decode())
    if fn := cursor.field_name:
        yield ((*ctx, "field"), fn)

    if cursor.goto_first_child():
        i = 0
        while True:
            yield from json_stream(cursor, (*ctx, "_", i))
            if not cursor.goto_next_sibling():
                break
            i += 1

        yield ((*ctx, "_"),)
        assert cursor.goto_parent()

    yield (ctx,)


tree = reqs.parse(sys.stdin.read())
stream = json_stream(tree.walk())
stream_text = json.dumps(list(stream))
stream_json = check_output(["jq", "fromstream(.[])"], input=stream_text, text=True)

print(stream_json)