q3alique-codeflow/codeflow/analysis/ast_parser.py

from __future__ import annotations
from pathlib import Path
from tree_sitter import Tree, Node as TSNode
from tree_sitter_languages import get_language, get_parser
from typing import Any


SUPPORTED_LANGUAGES = {
    "python", "javascript", "typescript", "java", "kotlin", "go",
    "ruby", "php", "rust", "c", "cpp", "c_sharp", "swift",
    "bash", "lua", "perl", "scala", "r", "elixir", "erlang",
}

_PARSER_CACHE: dict[str, Any] = {}


def _get_cached_parser(language: str):
    if language not in _PARSER_CACHE:
        _PARSER_CACHE[language] = get_parser(language)
    return _PARSER_CACHE[language]


def parse_file(file_path: str, language: str) -> Tree | None:
    if language == "unknown" or language not in SUPPORTED_LANGUAGES:
        return None
    try:
        source_bytes = Path(file_path).read_bytes()
    except (OSError, IOError):
        return None
    try:
        parser = _get_cached_parser(language)
        tree = parser.parse(source_bytes)
        return tree
    except Exception:
        return None


def query_tree(tree: Tree, language: str, query_string: str, capture_filter: str | set[str] | None = None) -> list[tuple[TSNode, str]]:
    try:
        lang = get_language(language)
        query = lang.query(query_string)
        captures = query.captures(tree.root_node)
        if capture_filter is not None:
            if isinstance(capture_filter, str):
                return [(n, c) for n, c in captures if c == capture_filter]
            elif isinstance(capture_filter, (set, list, tuple)):
                return [(n, c) for n, c in captures if c in capture_filter]
        return captures
    except Exception:
        return []


def get_node_text(ts_node: TSNode, source_bytes: bytes) -> str:
    try:
        text = source_bytes[ts_node.start_byte:ts_node.end_byte]
        return text.decode("utf-8", errors="replace").strip()
    except Exception:
        return ""