[docs]classSQLSegmenter(TreeSitterSegmenter):"""Code segmenter for SQL. This class uses Tree-sitter to segment SQL code into its constituent statements (e.g., SELECT, CREATE TABLE). It also provides functionality to extract these statements and simplify the code into commented descriptions. """
[docs]defget_language(self)->"Language":"""Return the SQL language grammar for Tree-sitter."""fromtree_sitter_languagesimportget_languagereturnget_language("sql")
[docs]defget_chunk_query(self)->str:"""Return the Tree-sitter query for SQL segmentation."""returnCHUNK_QUERY
[docs]defextract_functions_classes(self)->list[str]:"""Extract SQL statements from the code. Ensures that all SQL statements end with a semicolon for consistency. """extracted=super().extract_functions_classes()# Ensure all statements end with a semicolonreturn[stmt.strip()+";"ifnotstmt.strip().endswith(";")elsestmt.strip()forstmtinextracted]
[docs]defsimplify_code(self)->str:"""Simplify the extracted SQL code into comments. Converts SQL statements into commented descriptions for easy readability. """return"\n".join([f"-- Code for: {stmt.strip()}"forstmtinself.extract_functions_classes()])
[docs]defmake_line_comment(self,text:str)->str:"""Create a line comment in SQL style."""returnf"-- {text}"