Source code for langchain_community.document_loaders.parsers.language.tree_sitter_segmenter

from abc import abstractmethod
from typing import TYPE_CHECKING, List

from langchain_community.document_loaders.parsers.language.code_segmenter import (
    CodeSegmenter,
)

if TYPE_CHECKING:
    from tree_sitter import Language, Parser


class TreeSitterSegmenter(CodeSegmenter):
    """Abstract class for `CodeSegmenter`s that use the tree-sitter library."""

    def __init__(self, code: str):
        super().__init__(code)
        self.source_lines = self.code.splitlines()

        try:
            import tree_sitter  # noqa: F401
            import tree_sitter_languages  # noqa: F401
        except ImportError:
            raise ImportError(
                "Could not import tree_sitter/tree_sitter_languages Python packages. "
                "Please install them with "
                "`pip install tree-sitter tree-sitter-languages`."
            )

    def is_valid(self) -> bool:
        language = self.get_language()
        error_query = language.query("(ERROR) @error")

        parser = self.get_parser()
        tree = parser.parse(bytes(self.code, encoding="UTF-8"))

        return len(error_query.captures(tree.root_node)) == 0

    def extract_functions_classes(self) -> List[str]:
        language = self.get_language()
        query = language.query(self.get_chunk_query())

        parser = self.get_parser()
        tree = parser.parse(bytes(self.code, encoding="UTF-8"))
        captures = query.captures(tree.root_node)

        processed_lines = set()
        chunks = []

        for node, name in captures:
            start_line = node.start_point[0]
            end_line = node.end_point[0]
            lines = list(range(start_line, end_line + 1))

            if any(line in processed_lines for line in lines):
                continue

            processed_lines.update(lines)
            chunk_text = node.text.decode("UTF-8")
            chunks.append(chunk_text)

        return chunks

    def simplify_code(self) -> str:
        language = self.get_language()
        query = language.query(self.get_chunk_query())

        parser = self.get_parser()
        tree = parser.parse(bytes(self.code, encoding="UTF-8"))
        processed_lines = set()

        simplified_lines = self.source_lines[:]
        for node, name in query.captures(tree.root_node):
            start_line = node.start_point[0]
            end_line = node.end_point[0]

            lines = list(range(start_line, end_line + 1))
            if any(line in processed_lines for line in lines):
                continue

            simplified_lines[start_line] = self.make_line_comment(
                f"Code for: {self.source_lines[start_line]}"
            )

            for line_num in range(start_line + 1, end_line + 1):
                simplified_lines[line_num] = None  # type: ignore

            processed_lines.update(lines)

        return "\n".join(line for line in simplified_lines if line is not None)

    def get_parser(self) -> "Parser":
        from tree_sitter import Parser

        parser = Parser()
        parser.set_language(self.get_language())
        return parser

    @abstractmethod
    def get_language(self) -> "Language":
        raise NotImplementedError()  # pragma: no cover

    @abstractmethod
    def get_chunk_query(self) -> str:
        raise NotImplementedError()  # pragma: no cover

    @abstractmethod
    def make_line_comment(self, text: str) -> str:
        raise NotImplementedError()  # pragma: no cover