diff --git a/src/chonkie/chunker/code.py b/src/chonkie/chunker/code.py index 8144a8dda..48fe6ea71 100644 --- a/src/chonkie/chunker/code.py +++ b/src/chonkie/chunker/code.py @@ -308,14 +308,20 @@ def _create_chunks(self, """Create Code Chunks.""" chunks = [] current_index = 0 - for i in range(len(texts)): - text = texts[i] - token_count = token_counts[i] + for i, (text, token_count, group) in enumerate(zip(texts, token_counts, node_groups)): + start_line = None + end_line = None + if group: + # tree-sitter nodes provide 0-indexed row numbers, so add 1 for 1-indexed line numbers. + start_line = group[0].start_point[0] + 1 + end_line = group[-1].end_point[0] + 1 chunks.append(Chunk(text=text, start_index=current_index, end_index=current_index + len(text), - token_count=token_count)) # type: ignore[attr-defined] + token_count=token_count, + start_line=start_line, + end_line=end_line)) current_index += len(text) return chunks @@ -338,23 +344,33 @@ def chunk(self, text: str) -> List[Chunk]: else: logger.debug(f"Using configured language: {self.language}") + # Initialize node_groups here to ensure it's available in the finally block scope + node_groups = [] + chunks = [] try: # Create the parsing tree for the current code tree: Tree = self.parser.parse(original_text_bytes) # type: ignore root_node: Node = tree.root_node # type: ignore - # Get the node_groups + # Get the node_groups node_groups, token_counts = self._group_child_nodes(root_node) texts: List[str] = self._get_texts_from_node_groups(node_groups, original_text_bytes) - finally: + + # Create chunks inside the try block to ensure node_groups is available + chunks = self._create_chunks(texts, token_counts, node_groups) + logger.info(f"Created {len(chunks)} code chunks from parsed syntax tree") + + finally: # Clean up the tree and root_node if they are not needed if not self.include_nodes: - del tree, root_node - node_groups = [] + if 'tree' in locals(): + del tree + if 'root_node' in locals(): + del root_node + # Clear node_groups only after it's been used by _create_chunks + node_groups.clear() - chunks = self._create_chunks(texts, token_counts, node_groups) - logger.info(f"Created {len(chunks)} code chunks from parsed syntax tree") - return chunks + return chunks def __repr__(self) -> str: """Return the string representation of the CodeChunker.""" diff --git a/src/chonkie/types/base.py b/src/chonkie/types/base.py index 2badb175e..4a14a9d03 100644 --- a/src/chonkie/types/base.py +++ b/src/chonkie/types/base.py @@ -22,6 +22,8 @@ class Chunk: start_index (int): The starting index of the chunk in the original text. end_index (int): The ending index of the chunk in the original text. token_count (int): The number of tokens in the chunk. + start_line (Optional[int]): The starting line number of the chunk in the original text (1-indexed). + end_line (Optional[int]): The ending line number of the chunk in the original text (1-indexed). context (Optional[str]): Optional context metadata for the chunk. embedding (Union[List[float], np.ndarray, None]): Optional embedding vector for the chunk, either as a list of floats or a numpy array. @@ -33,6 +35,8 @@ class Chunk: start_index: int = field(default=0) end_index: int = field(default=0) token_count: int = field(default=0) + start_line: Optional[int] = field(default=None) + end_line: Optional[int] = field(default=None) context: Optional[str] = field(default=None) embedding: Union[List[float], "np.ndarray", None] = field(default=None) @@ -80,15 +84,17 @@ def _preview_embedding(self) -> str: def __repr__(self) -> str: """Return a detailed string representation of the Chunk.""" - repr = ( - f"Chunk(text='{self.text}', token_count={self.token_count}, " + repr_str = ( + f"Chunk(text='{self.text[:50]}...', token_count={self.token_count}, " f"start_index={self.start_index}, end_index={self.end_index}" ) + if self.start_line is not None and self.end_line is not None: + repr_str += f", start_line={self.start_line}, end_line={self.end_line}" if self.context: - repr += f", context={self.context}" + repr_str += f", context={self.context}" if self.embedding is not None: - repr += f", embedding={self._preview_embedding()}" - return repr + ")" + repr_str += f", embedding={self._preview_embedding()}" + return repr_str + ")" def __iter__(self) -> Iterator[str]: """Return an iterator over the chunk's text.""" @@ -119,6 +125,8 @@ def from_dict(cls, data: dict) -> "Chunk": start_index=data["start_index"], end_index=data["end_index"], token_count=data["token_count"], + start_line=data.get("start_line", None), + end_line=data.get("end_line", None), context=data.get("context", None), embedding=data.get("embedding", None), ) diff --git a/tests/chunkers/test_code_chunker.py b/tests/chunkers/test_code_chunker.py index 887702ad3..8a17ff72e 100644 --- a/tests/chunkers/test_code_chunker.py +++ b/tests/chunkers/test_code_chunker.py @@ -165,4 +165,27 @@ def test_code_chunker_chunk_size_javascript(js_code: str) -> None: chunks = chunker.chunk(js_code) # Allow for some leeway assert all(chunk.token_count < chunk_size + 15 for chunk in chunks[:-1]) - assert chunks[-1].token_count > 0 \ No newline at end of file + assert chunks[-1].token_count > 0 + + +def test_code_chunker_adds_line_numbers(python_code: str) -> None: + """Test that the CodeChunker correctly adds start and end line numbers to chunks.""" + # Use a chunk size that will reliably split the python_code fixture + chunker = CodeChunker(language="python", chunk_size=100) + chunks = chunker.chunk(python_code) + + assert len(chunks) > 0, "Chunker should produce at least one chunk" + + # Verify that every chunk has valid, 1-indexed line numbers + for chunk in chunks: + assert chunk.start_line is not None, "start_line should not be None" + assert chunk.end_line is not None, "end_line should not be None" + assert isinstance(chunk.start_line, int) + assert isinstance(chunk.end_line, int) + assert chunk.start_line > 0, "Line numbers should be 1-indexed and positive" + assert chunk.end_line >= chunk.start_line, "End line must be greater than or equal to start line" + + # The python_code fixture starts with a newline, so its content begins on line 2. + # The code content ends on line 19 of the input string. + assert chunks[0].start_line == 2, "The first chunk should start on the first line of code (line 2)" + assert chunks[-1].end_line == 19, "The last chunk should end on the last line of code (line 19)" \ No newline at end of file