fix(cli): preserve spaces and newlines in treesitter chunking.

Author: Zhe Yu

src/vectorcode/chunking.py

@@ -150,7 +150,7 @@ def __chunk_node(
                 f"Traversing at node {node.text.decode()} at position {node.byte_range}"
             )
         current_chunk: str = ""
-
+        prev_node = None
         current_start = None
 
         for child in node.children:
@@ -184,10 +184,19 @@ def __chunk_node(
                 current_start = Point(
                     row=child.start_point.row + 1, column=child.start_point.column
                 )
+                prev_node = child
 
-            elif len(current_chunk) + child_length <= self.config.chunk_size:
+            elif len(current_chunk) + child_length + 1 <= self.config.chunk_size:
                 # Add to current chunk
+                if prev_node:
+                    if prev_node.end_point.row != child.start_point.row:
+                        current_chunk += "\n"
+                    else:
+                        current_chunk += " " * (
+                            child.start_point.column - prev_node.end_point.column
+                        )
                 current_chunk += child_bytes.decode()
+                prev_node = child
 
             else:
                 # Yield current chunk and start new one

tests/test_chunking.py

@@ -250,7 +250,7 @@ def bar():
 
 
 def test_treesitter_chunker_filter_wildcard():
-    chunker = TreeSitterChunker(Config(chunk_size=30, chunk_filters={"*": [".*foo.*"]}))
+    chunker = TreeSitterChunker(Config(chunk_size=35, chunk_filters={"*": [".*foo.*"]}))
 
     test_content = r"""
 def foo():
@@ -283,12 +283,12 @@ def bar():
         test_file = tmp_file.name
 
     chunks = list(str(i) for i in chunker.chunk(test_file))
-    assert chunks == ['functionbar()return "bar"end']
+    assert chunks == ['function bar()\n  return "bar"\nend']
     os.remove(test_file)
 
 
 def test_treesitter_chunker_lua():
-    chunker = TreeSitterChunker(Config(chunk_size=30))
+    chunker = TreeSitterChunker(Config(chunk_size=35))
     test_content = r"""
 function foo()
   return "foo"
@@ -304,7 +304,10 @@ def test_treesitter_chunker_lua():
         test_file = tmp_file.name
 
     chunks = list(str(i) for i in chunker.chunk(test_file))
-    assert chunks == ['functionfoo()return "foo"end', 'functionbar()return "bar"end']
+    assert chunks == [
+        'function foo()\n  return "foo"\nend',
+        'function bar()\n  return "bar"\nend',
+    ]
 
     os.remove(test_file)
 
@@ -403,7 +406,7 @@ def bar():
     assert len(chunks) >= 2  # Should have at least 2 chunks
 
     # First chunk should contain the function definition start
-    assert "deffoo():" in chunks[0].text
+    assert "def foo():" in chunks[0].text
     assert chunks[0].start == Point(1, 0)
 
     # Last chunk should contain the final return statement