diff --git a/Cargo.lock b/Cargo.lock index 7a85d1dac..2226da8ef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1978,6 +1978,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -12002,6 +12003,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/docs/src/content/docs/ops/text.mdx b/docs/src/content/docs/ops/text.mdx index 1f25a4c58..63e9952a4 100644 --- a/docs/src/content/docs/ops/text.mdx +++ b/docs/src/content/docs/ops/text.mdx @@ -155,6 +155,7 @@ Languages with syntax-aware (tree-sitter) splitting — splits at logical bounda | JSON | `"json"` | `.json`, `.jsonl`, `.geojson` | | Julia | `"julia"` | `.jl` | | Kotlin | `"kotlin"` | `.kt`, `.kts`, `.ktm` | +| Lua | `"lua"` | `.lua`, `.nse`, `.p8`, `.pd_lua`, `.rbxs`, `.rockspec`, `.wlua` | | Markdown | `"markdown"` | `.md`, `.markdown`, `.mdx` | | Pascal / Delphi | `"pascal"` | `.pas`, `.dpr`, `.lpr`, `.dfm` | | PHP | `"php"` | `.php` | @@ -176,7 +177,7 @@ Languages with syntax-aware (tree-sitter) splitting — splits at logical bounda Note that `.h` maps to C++, not C — the C++ grammar also parses C headers. Pass `language="c"` explicitly to force the C grammar. -Over 110 additional languages use separator-based splitting (e.g. `"elixir"`, `"erlang"`, `"haskell"`, `"lua"`, `"nix"`, `"perl"`, `"powershell"`, `"proto"`, `"zig"`, and more). A `language=` value that matches no known language also falls back to separator-based splitting rather than raising. +110 additional languages use separator-based splitting (e.g. `"elixir"`, `"erlang"`, `"haskell"`, `"nix"`, `"perl"`, `"powershell"`, `"proto"`, `"zig"`, and more). A `language=` value that matches no known language also falls back to separator-based splitting rather than raising. ### `CustomLanguageConfig` diff --git a/examples/rust/amazon_s3_embedding/Cargo.lock b/examples/rust/amazon_s3_embedding/Cargo.lock index b03c4c7c5..7173b25ce 100644 --- a/examples/rust/amazon_s3_embedding/Cargo.lock +++ b/examples/rust/amazon_s3_embedding/Cargo.lock @@ -1085,6 +1085,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -5631,6 +5632,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/code_embedding/Cargo.lock b/examples/rust/code_embedding/Cargo.lock index 37a7023bc..8c06c5349 100644 --- a/examples/rust/code_embedding/Cargo.lock +++ b/examples/rust/code_embedding/Cargo.lock @@ -529,6 +529,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -4670,6 +4671,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/code_embedding_lancedb/Cargo.lock b/examples/rust/code_embedding_lancedb/Cargo.lock index af8330650..569a9f94e 100644 --- a/examples/rust/code_embedding_lancedb/Cargo.lock +++ b/examples/rust/code_embedding_lancedb/Cargo.lock @@ -955,6 +955,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -7270,6 +7271,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/gdrive_text_embedding/Cargo.lock b/examples/rust/gdrive_text_embedding/Cargo.lock index a1963e8e4..500143abf 100644 --- a/examples/rust/gdrive_text_embedding/Cargo.lock +++ b/examples/rust/gdrive_text_embedding/Cargo.lock @@ -533,6 +533,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -4675,6 +4676,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/oci_object_storage_embedding/Cargo.lock b/examples/rust/oci_object_storage_embedding/Cargo.lock index 98c484510..41d48fa8c 100644 --- a/examples/rust/oci_object_storage_embedding/Cargo.lock +++ b/examples/rust/oci_object_storage_embedding/Cargo.lock @@ -533,6 +533,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -4675,6 +4676,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/paper_metadata/Cargo.lock b/examples/rust/paper_metadata/Cargo.lock index e35d07ae5..f565ed058 100644 --- a/examples/rust/paper_metadata/Cargo.lock +++ b/examples/rust/paper_metadata/Cargo.lock @@ -574,6 +574,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -4786,6 +4787,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/pdf_embedding/Cargo.lock b/examples/rust/pdf_embedding/Cargo.lock index 208270235..e2e132f19 100644 --- a/examples/rust/pdf_embedding/Cargo.lock +++ b/examples/rust/pdf_embedding/Cargo.lock @@ -574,6 +574,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -4784,6 +4785,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/text_embedding/Cargo.lock b/examples/rust/text_embedding/Cargo.lock index 44ae52b1f..82b257c57 100644 --- a/examples/rust/text_embedding/Cargo.lock +++ b/examples/rust/text_embedding/Cargo.lock @@ -529,6 +529,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -4670,6 +4671,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/text_embedding_lancedb/Cargo.lock b/examples/rust/text_embedding_lancedb/Cargo.lock index d96d19ce5..3940e9eda 100644 --- a/examples/rust/text_embedding_lancedb/Cargo.lock +++ b/examples/rust/text_embedding_lancedb/Cargo.lock @@ -955,6 +955,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -7270,6 +7271,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/text_embedding_qdrant/Cargo.lock b/examples/rust/text_embedding_qdrant/Cargo.lock index c10819e70..73e287462 100644 --- a/examples/rust/text_embedding_qdrant/Cargo.lock +++ b/examples/rust/text_embedding_qdrant/Cargo.lock @@ -545,6 +545,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -4478,6 +4479,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/examples/rust/text_embedding_turbopuffer/Cargo.lock b/examples/rust/text_embedding_turbopuffer/Cargo.lock index 29bca684e..75d473258 100644 --- a/examples/rust/text_embedding_turbopuffer/Cargo.lock +++ b/examples/rust/text_embedding_turbopuffer/Cargo.lock @@ -513,6 +513,7 @@ dependencies = [ "tree-sitter-julia", "tree-sitter-kotlin-ng", "tree-sitter-language", + "tree-sitter-lua", "tree-sitter-md", "tree-sitter-pascal", "tree-sitter-php", @@ -4129,6 +4130,16 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" +[[package]] +name = "tree-sitter-lua" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8daaf5f4235188a58603c39760d5fa5d4b920d36a299c934adddae757f32a10c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-md" version = "0.5.3" diff --git a/python/tests/ops/test_text.py b/python/tests/ops/test_text.py index 6037f8f4e..bbaf77de3 100644 --- a/python/tests/ops/test_text.py +++ b/python/tests/ops/test_text.py @@ -19,6 +19,7 @@ def test_detect_code_language_known_extensions() -> None: assert detect_code_language(filename="App.vue") == "vue" assert detect_code_language(filename="script.jl") == "julia" assert detect_code_language(filename="main.dart") == "dart" + assert detect_code_language(filename="init.lua") == "lua" assert detect_code_language(filename="Main.elm") == "elm" assert detect_code_language(filename="index.astro") == "astro" assert detect_code_language(filename="deploy.sh") == "bash" @@ -229,6 +230,54 @@ def test_recursive_splitter_with_dart() -> None: assert all(isinstance(c, Chunk) for c in chunks) +def test_recursive_splitter_with_lua() -> None: + """Lua splitting is syntax-aware, not separator-based. + + The sample is one line, so separators give the fallback path no function + boundaries at all. Only the grammar can find them. + """ + splitter = RecursiveSplitter() + code = ( + "local function f(a) return a+1 end " + "local function g(b) return b*2 end " + "local function h(c) return c-3 end" + ) + chunks = splitter.split(code, chunk_size=45, min_chunk_size=10, language="lua") + + assert len(chunks) > 1, "sample must split for the check below to mean anything" + assert all(isinstance(c, Chunk) for c in chunks) + for chunk in chunks: + text = chunk.text.strip() + assert text.startswith("local function") and text.endswith("end"), ( + f"chunk straddles a function boundary: {text!r}" + ) + + +def test_recursive_splitter_lua_keeps_table_constructor_whole() -> None: + """A Lua table constructor is not severed at a comma between its fields. + + Separator fallback cuts `{ x = x, y = y }` at the comma, because a comma is + a separator to it and a table field to the grammar. A constructor longer + than `chunk_size` is still split, here as in every other grammar-backed + language. + """ + splitter = RecursiveSplitter() + code = ( + "local function clamp(v, lo, hi) return math.min(math.max(v, lo), hi) end\n" + "local Point = {}\n" + "Point.__index = Point\n" + "function Point.new(x, y) return setmetatable({ x = x, y = y }, Point) end\n" + "function Point:norm() return math.sqrt(self.x * self.x + self.y * self.y) end\n" + "return Point\n" + ) + chunks = splitter.split(code, chunk_size=60, min_chunk_size=10, language="lua") + + assert len(chunks) > 1, "sample must split for the check below to mean anything" + assert any("setmetatable({ x = x, y = y }, Point)" in c.text for c in chunks), ( + "table constructor was severed across chunks" + ) + + def test_recursive_splitter_with_vue() -> None: """Test RecursiveSplitter with Vue syntax-aware splitting.""" splitter = RecursiveSplitter() diff --git a/rust/code_ast/Cargo.toml b/rust/code_ast/Cargo.toml index 4b1a778a1..0353b3d8f 100644 --- a/rust/code_ast/Cargo.toml +++ b/rust/code_ast/Cargo.toml @@ -30,6 +30,7 @@ tree-sitter-json = "0.24.8" tree-sitter-julia = "0.23.1" # The other more popular crate tree-sitter-kotlin requires tree-sitter < 0.23 for now tree-sitter-kotlin-ng = "1.1.0" +tree-sitter-lua = "0.5.0" tree-sitter-md = "0.5.3" tree-sitter-pascal = "0.10.2" tree-sitter-php = "0.23.11" diff --git a/rust/code_ast/src/prog_langs.rs b/rust/code_ast/src/prog_langs.rs index 13dd4ea18..66a1c9e1f 100644 --- a/rust/code_ast/src/prog_langs.rs +++ b/rust/code_ast/src/prog_langs.rs @@ -328,7 +328,7 @@ static LANGUAGE_INFO_BY_NAME: LazyLock, &'static Program ".rockspec", ".wlua", ], - None, + Some(TreeSitterLanguageInfo::new(tree_sitter_lua::LANGUAGE, [])), ); add("luau", &[".luau"], None); add("magik", &[".magik"], None); @@ -622,6 +622,13 @@ mod tests { assert!(dart.treesitter_info.is_some()); } + #[test] + fn test_lua_has_treesitter() { + let lua = get_language_info(".lua").unwrap(); + assert_eq!(lua.name.as_ref(), "lua"); + assert!(lua.treesitter_info.is_some()); + } + #[test] fn test_astro_has_treesitter() { let astro = get_language_info(".astro").unwrap(); diff --git a/rust/ops_text/src/split/recursive.rs b/rust/ops_text/src/split/recursive.rs index 741a825a0..f1e80240e 100644 --- a/rust/ops_text/src/split/recursive.rs +++ b/rust/ops_text/src/split/recursive.rs @@ -932,6 +932,64 @@ class Point { assert!(!chunks.is_empty()); } + #[test] + fn test_split_with_lua_language() { + let chunker = RecursiveChunker::new(RecursiveSplitConfig::default()).unwrap(); + // One line, so line/whitespace separators give the fallback path no + // function boundaries at all. Only the grammar can find them. + let text = "local function f(a) return a+1 end \ + local function g(b) return b*2 end \ + local function h(c) return c-3 end"; + let config = RecursiveChunkConfig { + chunk_size: 45, + min_chunk_size: Some(10), + chunk_overlap: Some(0), + }; + let chunks = chunker.split(&CodeSource::with_language(text, "lua"), config); + assert!( + chunks.len() > 1, + "sample must split for the check below to mean anything" + ); + for chunk in &chunks { + let chunk_text = text[chunk.range.start..chunk.range.end].trim(); + assert!( + chunk_text.starts_with("local function") && chunk_text.ends_with("end"), + "chunk straddles a function boundary: {chunk_text:?}" + ); + } + } + + #[test] + fn test_split_lua_keeps_table_constructor_whole() { + let chunker = RecursiveChunker::new(RecursiveSplitConfig::default()).unwrap(); + // Asserts one thing: the argument list of `setmetatable` stays whole. + // Separator fallback cuts it at the comma between `x = x` and `y = y`, + // because a comma is a separator to it and a table field to the + // grammar. A constructor longer than `chunk_size` is still split, here + // as in every other grammar-backed language. + let text = "local function clamp(v, lo, hi) return math.min(math.max(v, lo), hi) end\n\ + local Point = {}\n\ + Point.__index = Point\n\ + function Point.new(x, y) return setmetatable({ x = x, y = y }, Point) end\n\ + function Point:norm() return math.sqrt(self.x * self.x + self.y * self.y) end\n\ + return Point\n"; + let config = RecursiveChunkConfig { + chunk_size: 60, + min_chunk_size: Some(10), + chunk_overlap: Some(0), + }; + let chunks = chunker.split(&CodeSource::with_language(text, "lua"), config); + assert!( + chunks.len() > 1, + "sample must split for the check below to mean anything" + ); + assert!( + chunks.iter().any(|c| text[c.range.start..c.range.end] + .contains("setmetatable({ x = x, y = y }, Point)")), + "table constructor was severed across chunks" + ); + } + #[test] fn test_split_positions() { let chunker = RecursiveChunker::new(RecursiveSplitConfig::default()).unwrap();