ASTChunk: AST-based Code Chunking Library

You can install ASTChunk via PyPI or from source. The package automatically installs its core dependencies, including tree-sitter and the necessary language parsers for Python, Java, C#, and TypeScript.

# From PyPI
pip install astchunk

# From source
git clone git@github.com:yilinjz/astchunk.git
pip install -e .

When initializing ASTChunkBuilder or calling .chunkify(), you can use the following configuration keys:

• max_chunk_size: Maximum non-whitespace characters per chunk. Note that the final size after overlapping or expansion may exceed this limit.
• language: The programming language for parsing (python, java, csharp, or typescript).
• metadata_template: The format for the output metadata (e.g., "default").
• repo_level_metadata (optional): A dictionary containing repository-wide information like filepath.
• chunk_overlap (optional): The number of AST nodes to overlap between consecutive chunks.
• chunk_expansion (optional): A boolean indicating whether to add metadata headers to the chunks.

You can enhance chunking by enabling node overlapping and metadata expansion. Overlapping helps maintain context between chunks, while expansion adds metadata headers directly into the chunk content. These can be set globally in the builder or passed as arguments to chunkify() for single-use configurations.

# Example of advanced configuration
configs = {
    "max_chunk_size": 100,
    "language": "python",
    "metadata_template": "default",
    "repo_level_metadata": {
        "filepath": "src/calculator.py"
    },
    "chunk_overlap": 1,
    "chunk_expansion": True
}

chunk_builder = ASTChunkBuilder(**configs)
chunks = chunk_builder.chunkify(code)

Create separate builder instances for different languages to ensure correct AST parsing.

# Python code
python_builder = ASTChunkBuilder(
    max_chunk_size=1500,
    language="python",
    metadata_template="default"
)

# Java code  
java_builder = ASTChunkBuilder(
    max_chunk_size=2000,
    language="java", 
    metadata_template="default"
)

# TypeScript code
ts_builder = ASTChunkBuilder(
    max_chunk_size=1800,
    language="typescript",
    metadata_template="default"
)

The build_metadata method populates the chunk's metadata based on the metadata_template provided during initialization. This allows the chunk to be compatible with different evaluation benchmarks or general use cases.

Supported Templates:

• "none": Creates an empty metadata dictionary.
• "default": Includes filepath, chunk_size, line_count, start_line_no, end_line_no, and node_count. Requires repo_level_metadata to contain "filepath".
• "coderagbench-repoeval": Includes fpath_tuple, repo, chunk_size, line_count, start_line_no, end_line_no, and node_count. Requires repo_level_metadata to contain "fpath_tuple" and "repo".
• "coderagbench-swebench-lite": Includes _id (formatted as {instance_id}_{start_line}-{end_line}) and title. Requires repo_level_metadata to contain "instance_id" and "filename".

ASTChunk provides full support for the following languages via tree-sitter parsers:

The preprocessing module provides low-level utilities for analyzing source code structure, specifically focusing on whitespace (NWS - Non-Whitespace characters) and byte/integer ranges. These are useful for fine-grained control over how chunks are split or measured.

Key utilities include:

• preprocess_nws_count: Pre-calculates whitespace counts for efficient processing.
• get_nws_count / get_nws_count_direct: Retrieves whitespace counts.
• get_nodes_in_brange: Finds nodes within a specific ByteRange.
• get_largest_node_in_brange: Identifies the most significant node within a given range.

from astchunk import (
    ByteRange, 
    IntRange, 
    preprocess_nws_count,
    get_nodes_in_brange
)

# Example usage of range types
brange_obj = ByteRange(start=0, end=100)
# Example usage of preprocessing
# preprocess_nws_count(source_code)

The ASTChunkBuilder class is the primary entry point for parsing code into structurally-aware chunks using Abstract Syntax Trees (AST). When initializing, you must provide configuration for the maximum chunk size, the programming language, and a metadata template.

Supported languages include:

• python
• java
• csharp
• typescript

from astchunk.astchunk_builder import ASTChunkBuilder

builder = ASTChunkBuilder(
    max_chunk_size=500,
    language="python",
    metadata_template="standard"
)

Use get_nodes_in_brange to retrieve all valid tree-sitter nodes that are fully contained within a specific ByteRange.

• The function traverses the syntax tree starting from a provided root_node.
• It excludes nodes of type "ERROR" and their descendants.
• It handles "module" nodes by traversing their children.

import tree_sitter as ts
from astchunk.preprocessing import get_nodes_in_brange, ByteRange

# Assuming root_node is a valid tree-sitter node
my_brange = ByteRange(start=10, stop=50)
nodes = get_nodes_in_brange(root_node, my_brange)

for node in nodes:
    print(f"Node type: {node.type}, Range: {node.start_byte}-{node.end_byte}")

The get_largest_node_in_brange function returns the size of the largest node found within a specified ByteRange. You can choose how size is measured using the size_option parameter:

• "byte": Measures size in raw bytes (end_byte - start_byte).
• "non-ws": Measures size based on the count of non-whitespace characters.

If no nodes are found within the range, it returns 0.

from astchunk.preprocessing import get_largest_node_in_brange, ByteRange

# Get largest node size in bytes
byte_size = get_largest_node_in_brange(ts_node, ByteRange(0, 100), size_option="byte")

# Get largest node size in non-whitespace characters
char_size = get_largest_node_in_brange(ts_node, ByteRange(0, 100), size_option="non-ws")

The to_code_window method converts the ASTChunk object into a dictionary format suitable for downstream integration (e.g., feeding into an LLM or a search engine).

Output Formats:

• If metadata_template is "coderagbench-swebench-lite": Returns a dict with _id, title, and text keys.

  {
    "_id": "...",
    "title": "...",
    "text": "..."
  }

• Otherwise: Returns a dict with content and metadata keys.

  {
    "content": "...",
    "metadata": { ... }
  }