CintraAI/code-chunker
5
1import os2import subprocess3from typing import List, Dict, Union, Tuple4from tree_sitter import Language, Parser, Node5import logging6 7class CodeParser:8 # Added a CACHE_DIR class attribute for caching9 CACHE_DIR = os.path.expanduser("~/.code_parser_cache")10 11 def __init__(self, file_extensions: Union[None, List[str], str] = None):12 if isinstance(file_extensions, str):13 file_extensions = [file_extensions]14 self.language_extension_map = {15 "py": "python",16 "js": "javascript",17 "jsx": "javascript",18 "css": "css",19 "ts": "typescript",20 "tsx": "typescript",21 "php": "php",22 "rb": "ruby",23 "go": "go"24 }25 if file_extensions is None:26 self.language_names = []27 else:28 self.language_names = [self.language_extension_map.get(ext) for ext in file_extensions if29 ext in self.language_extension_map]30 self.languages = {}31 self._install_parsers()32 33 def _install_parsers(self):34 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')35 36 try:37 # Ensure cache directory exists38 if not os.path.exists(self.CACHE_DIR):39 os.makedirs(self.CACHE_DIR)40 41 for language in self.language_names:42 repo_path = os.path.join(self.CACHE_DIR, f"tree-sitter-{language}")43 44 # Check if the repository exists and contains necessary files45 if not os.path.exists(repo_path) or not self._is_repo_valid(repo_path, language):46 try:47 if os.path.exists(repo_path):48 logging.info(f"Updating existing repository for {language}")49 update_command = f"cd {repo_path} && git pull"50 subprocess.run(update_command, shell=True, check=True)51 else:52 logging.info(f"Cloning repository for {language}")53 clone_command = f"git clone https://github.com/tree-sitter/tree-sitter-{language} {repo_path}"54 subprocess.run(clone_command, shell=True, check=True)55 except subprocess.CalledProcessError as e:56 logging.error(f"Failed to clone/update repository for {language}. Error: {e}")57 continue58 59 try:60 build_path = os.path.join(self.CACHE_DIR, f"build/{language}.so")61 62 # Special handling for TypeScript63 if language == 'typescript':64 ts_dir = os.path.join(repo_path, 'typescript')65 tsx_dir = os.path.join(repo_path, 'tsx')66 if os.path.exists(ts_dir) and os.path.exists(tsx_dir):67 Language.build_library(build_path, [ts_dir, tsx_dir])68 else:69 raise FileNotFoundError(f"TypeScript or TSX directory not found in {repo_path}")70 elif language == 'php':71 php_dir = os.path.join(repo_path, 'php')72 if os.path.exists(php_dir):73 Language.build_library(build_path, [php_dir])74 else:75 raise FileNotFoundError(f"PHP directory not found in {repo_path}")76 else:77 Language.build_library(build_path, [repo_path])78 79 self.languages[language] = Language(build_path, language)80 logging.info(f"Successfully built and loaded {language} parser")81 except Exception as e:82 logging.error(f"Failed to build or load language {language}. Error: {str(e)}")83 logging.error(f"Repository path: {repo_path}")84 logging.error(f"Build path: {build_path}")85 if language == 'typescript':86 logging.error(f"TypeScript dir exists: {os.path.exists(ts_dir)}")87 logging.error(f"TSX dir exists: {os.path.exists(tsx_dir)}")88 elif language == 'php':89 logging.error(f"PHP dir exists: {os.path.exists(php_dir)}")90 91 except Exception as e:92 logging.error(f"An unexpected error occurred during parser installation: {str(e)}")93 94 def _is_repo_valid(self, repo_path: str, language: str) -> bool:95 """Check if the repository contains necessary files."""96 if language == 'typescript':97 return (os.path.exists(os.path.join(repo_path, 'typescript', 'src', 'parser.c')) and98 os.path.exists(os.path.join(repo_path, 'tsx', 'src', 'parser.c')))99 elif language == 'php':100 return os.path.exists(os.path.join(repo_path, 'php', 'src', 'parser.c'))101 else:102 return os.path.exists(os.path.join(repo_path, 'src', 'parser.c'))103 104 def parse_code(self, code: str, file_extension: str) -> Union[None, Node]:105 language_name = self.language_extension_map.get(file_extension)106 if language_name is None:107 print(f"Unsupported file type: {file_extension}")108 return None109 110 language = self.languages.get(language_name)111 if language is None:112 print("Language parser not found")113 return None114 115 parser = Parser()116 parser.set_language(language)117 tree = parser.parse(bytes(code, "utf8"))118 119 if tree is None:120 print("Failed to parse the code")121 return None122 123 return tree.root_node124 125 def extract_points_of_interest(self, node: Node, file_extension: str) -> List[Tuple[Node, str]]:126 node_types_of_interest = self._get_node_types_of_interest(file_extension)127 128 points_of_interest = []129 if node.type in node_types_of_interest.keys():130 points_of_interest.append((node, node_types_of_interest[node.type]))131 132 for child in node.children:133 points_of_interest.extend(self.extract_points_of_interest(child, file_extension))134 135 return points_of_interest136 137 def _get_node_types_of_interest(self, file_extension: str) -> Dict[str, str]:138 node_types = {139 'py': {140 'import_statement': 'Import',141 'export_statement': 'Export',142 'class_definition': 'Class',143 'function_definition': 'Function',144 },145 'css': {146 'tag_name': 'Tag',147 '@media': 'Media Query',148 },149 'js': {150 'import_statement': 'Import',151 'export_statement': 'Export',152 'class_declaration': 'Class',153 'function_declaration': 'Function',154 'arrow_function': 'Arrow Function',155 'statement_block': 'Block',156 },157 'ts': {158 'import_statement': 'Import',159 'export_statement': 'Export',160 'class_declaration': 'Class',161 'function_declaration': 'Function',162 'arrow_function': 'Arrow Function',163 'statement_block': 'Block',164 'interface_declaration': 'Interface',165 'type_alias_declaration': 'Type Alias',166 },167 'php': {168 'namespace_definition': 'Namespace',169 'class_declaration': 'Class',170 'method_declaration': 'Method',171 'function_definition': 'Function',172 'interface_declaration': 'Interface',173 'trait_declaration': 'Trait',174 },175 'rb': {176 'class': 'Class',177 'method': 'Method',178 'module': 'Module',179 'singleton_class': 'Singleton Class',180 'begin': 'Begin Block',181 },182 'go': {183 'import_declaration': 'Import',184 'function_declaration': 'Function',185 'method_declaration': 'Method',186 'type_declaration': 'Type',187 'struct_type': 'Struct',188 'interface_type': 'Interface',189 'package_clause': 'Package'190 }191 }192 193 if file_extension in node_types.keys():194 return node_types[file_extension]195 elif file_extension == "jsx":196 return node_types["js"]197 elif file_extension == "tsx":198 return node_types["ts"]199 else:200 raise ValueError("Unsupported file type")201 202 203 def _get_nodes_for_comments(self, file_extension: str) -> Dict[str, str]:204 node_types = {205 'py': {206 'comment': 'Comment',207 'decorator': 'Decorator', # Broadened category208 },209 'css': {210 'comment': 'Comment'211 },212 'js': {213 'comment': 'Comment',214 'decorator': 'Decorator', # Broadened category215 },216 'ts': {217 'comment': 'Comment',218 'decorator': 'Decorator',219 },220 'php': {221 'comment': 'Comment',222 'attribute': 'Attribute',223 },224 'rb': {225 'comment': 'Comment',226 },227 'go': {228 'comment': 'Comment',229 }230 }231 232 if file_extension in node_types.keys():233 return node_types[file_extension]234 elif file_extension == "jsx":235 return node_types["js"]236 elif file_extension == "tsx":237 return node_types["ts"]238 else:239 raise ValueError("Unsupported file type")240 241 def extract_comments(self, node: Node, file_extension: str) -> List[Tuple[Node, str]]:242 node_types_of_interest = self._get_nodes_for_comments(file_extension)243 244 comments = []245 if node.type in node_types_of_interest:246 comments.append((node, node_types_of_interest[node.type]))247 248 for child in node.children:249 comments.extend(self.extract_comments(child, file_extension))250 251 return comments252 253 def get_lines_for_points_of_interest(self, code: str, file_extension: str) -> List[int]:254 language_name = self.language_extension_map.get(file_extension)255 if language_name is None:256 raise ValueError("Unsupported file type")257 258 language = self.languages.get(language_name)259 if language is None:260 raise ValueError("Language parser not found")261 262 parser = Parser()263 parser.set_language(language)264 265 tree = parser.parse(bytes(code, "utf8"))266 267 root_node = tree.root_node268 points_of_interest = self.extract_points_of_interest(root_node, file_extension)269 270 line_numbers_with_type_of_interest = {}271 272 for node, type_of_interest in points_of_interest:273 start_line = node.start_point[0] 274 if type_of_interest not in line_numbers_with_type_of_interest:275 line_numbers_with_type_of_interest[type_of_interest] = []276 277 if start_line not in line_numbers_with_type_of_interest[type_of_interest]:278 line_numbers_with_type_of_interest[type_of_interest].append(start_line)279 280 lines_of_interest = []281 for _, line_numbers in line_numbers_with_type_of_interest.items():282 lines_of_interest.extend(line_numbers)283 284 return lines_of_interest285 286 def get_lines_for_comments(self, code: str, file_extension: str) -> List[int]:287 language_name = self.language_extension_map.get(file_extension)288 if language_name is None:289 raise ValueError("Unsupported file type")290 291 language = self.languages.get(language_name)292 if language is None:293 raise ValueError("Language parser not found")294 295 parser = Parser()296 parser.set_language(language)297 298 tree = parser.parse(bytes(code, "utf8"))299 300 root_node = tree.root_node301 comments = self.extract_comments(root_node, file_extension)302 303 line_numbers_with_comments = {}304 305 for node, type_of_interest in comments:306 start_line = node.start_point[0] 307 if type_of_interest not in line_numbers_with_comments:308 line_numbers_with_comments[type_of_interest] = []309 310 if start_line not in line_numbers_with_comments[type_of_interest]:311 line_numbers_with_comments[type_of_interest].append(start_line)312 313 lines_of_interest = []314 for _, line_numbers in line_numbers_with_comments.items():315 lines_of_interest.extend(line_numbers)316 317 return lines_of_interest318 319 def print_all_line_types(self, code: str, file_extension: str):320 language_name = self.language_extension_map.get(file_extension)321 if language_name is None:322 print(f"Unsupported file type: {file_extension}")323 return324 325 language = self.languages.get(language_name)326 if language is None:327 print("Language parser not found")328 return329 330 parser = Parser()331 parser.set_language(language)332 tree = parser.parse(bytes(code, "utf8"))333 334 root_node = tree.root_node335 line_to_node_type = self.map_line_to_node_type(root_node)336 337 code_lines = code.split('\n')338 339 for line_num, node_types in line_to_node_type.items():340 line_content = code_lines[line_num - 1] # Adjusting index for zero-based indexing341 print(f"line {line_num}: {', '.join(node_types)} | Code: {line_content}")342 343 344 def map_line_to_node_type(self, node, line_to_node_type=None, depth=0):345 if line_to_node_type is None:346 line_to_node_type = {}347 348 start_line = node.start_point[0] + 1 # Tree-sitter lines are 0-indexed; converting to 1-indexed349 350 # Only add the node type if it's the start line of the node351 if start_line not in line_to_node_type:352 line_to_node_type[start_line] = []353 line_to_node_type[start_line].append(node.type)354 355 for child in node.children:356 self.map_line_to_node_type(child, line_to_node_type, depth + 1)357 358 return line_to_node_type359