Phase 1: Rust/PyO3 acceleration for SnakemakeFormatter.parse()

.gitignore

@@ -135,3 +135,8 @@ dmypy.json
 
 # version
 _version.py
+
+# Rust build artifacts
+target/
+Cargo.lock
+

Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "snakebids-rust"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+name = "_core"
+path = "rust/src/lib.rs"
+crate-type = ["cdylib"]
+
+[dependencies]
+pyo3 = { version = "0.22", features = ["extension-module"] }

docs/general/rust_acceleration.md

@@ -0,0 +1,46 @@
+# Optional Rust Acceleration
+
+`snakebids` ships with an optional compiled extension
+(`snakebids._rust._core`) that accelerates `SnakemakeFormatter.parse()`.
+The extension is a private implementation detail—no public API changes
+are required to use it.
+
+## How it works
+
+When the extension is present, `SnakemakeFormatter.parse()` delegates
+its inner loop to a Rust implementation built with
+[PyO3](https://pyo3.rs/). The Python fallback is always available and is
+used automatically when the compiled extension is absent.
+
+## Building the extension locally
+
+You will need [Rust](https://rustup.rs/) and
+[maturin](https://www.maturin.rs/) installed.
+
+```bash
+# install maturin once
+pip install maturin
+
+# build and place the extension inside the source tree
+maturin develop --release
+```
+
+After this, `snakebids._rust._core` is importable and all
+`SnakemakeFormatter` calls benefit from the faster implementation
+transparently.
+
+## Verifying the extension is active
+
+```python
+from snakebids.utils.snakemake_templates import _HAS_RUST_PARSE
+print(_HAS_RUST_PARSE)  # True when the extension is loaded
+```
+
+## CI / packaging notes
+
+* Pure-Python installs (`pip install snakebids`) **do not** require
+  Rust—the extension is optional.
+* Pre-built wheels distributed on PyPI will include the compiled
+  extension for supported platforms.
+* The `[tool.maturin]` section in `pyproject.toml` configures the
+  module placement (`snakebids._rust._core`) and Python source root.

pyproject.toml

@@ -277,6 +277,10 @@ more_itertools = "itx"
 [tool.ruff.lint.flake8-pytest-style]
 fixture-parentheses = false
 
+[tool.maturin]
+module-name = "snakebids._rust._core"
+python-source = "src"
+
 [tool.codespell]
 # Ref: https://github.com/codespell-project/codespell#using-a-config-file
 skip = '.git*,*.lock,*.css,./typings'

rust/src/lib.rs

@@ -0,0 +1,159 @@
+//! Rust-accelerated internals for snakebids.
+//!
+//! This module exposes `parse_format_string`, a fast equivalent of the Python
+//! `SnakemakeFormatter.parse()` parsing loop.
+
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+
+const UNEXPECTED_CLOSE: &str = "unexpected '}' in string";
+const MISSING_CLOSE: &str = "expected '}' before end of string";
+const UNEXPECTED_OPEN: &str = "unexpected '{' in field name";
+
+/// Parse a Snakemake-style format string character by character.
+///
+/// Returns a list of 5-tuples per parsed segment:
+///   `(literal_text, field_name, new_underscore, new_current_field, update_current_field)`
+///
+/// - `literal_text`         – the literal portion before the next field (or at end of string)
+/// - `field_name`           – `None` for non-field segments; the field name for field segments
+/// - `new_underscore`       – `None` = no change; `Some(s)` = set `_underscore` to `s`
+/// - `new_current_field`    – new `_current_field` value (only meaningful when `update_current_field`)
+/// - `update_current_field` – whether `_current_field` should be updated
+///
+/// Raises `ValueError` on malformed input (same conditions as the Python implementation).
+#[pyfunction]
+pub fn parse_format_string(
+    format_string: &str,
+) -> PyResult<Vec<(String, Option<String>, Option<String>, Option<String>, bool)>> {
+    let mut entries: Vec<(String, Option<String>, Option<String>, Option<String>, bool)> =
+        Vec::new();
+
+    let mut chars = format_string.chars();
+
+    // Accumulates literal text between fields.
+    let mut literal = String::new();
+
+    loop {
+        match chars.next() {
+            // ---- End of string ------------------------------------------
+            None => {
+                if !literal.is_empty() {
+                    entries.push((literal, None, None, None, false));
+                }
+                return Ok(entries);
+            }
+
+            // ---- Opening brace ------------------------------------------
+            Some('{') => {
+                match chars.next() {
+                    None => {
+                        // Trailing lone `{` — no closing brace
+                        return Err(PyValueError::new_err(MISSING_CLOSE));
+                    }
+                    Some('{') => {
+                        // `{{` — escaped open brace (emits up-to-and-including first `{`)
+                        literal.push('{');
+                        entries.push((literal, None, Some("_".to_string()), None, false));
+                        literal = String::new();
+                    }
+                    Some('}') => {
+                        // `{}` — empty field name
+                        let new_underscore = literal.chars().last().map(|last| {
+                            if last == '/' || last == '_' {
+                                String::new()
+                            } else {
+                                "_".to_string()
+                            }
+                        });
+                        entries.push((
+                            literal,
+                            Some(String::new()),
+                            new_underscore,
+                            None,
+                            true,
+                        ));
+                        literal = String::new();
+                    }
+                    Some(first) => {
+                        // Start of a real field.  Collect everything up to the matching `}`.
+                        let mut field_content = String::new();
+                        field_content.push(first);
+
+                        let closing = loop {
+                            match chars.next() {
+                                None => {
+                                    return Err(PyValueError::new_err(MISSING_CLOSE));
+                                }
+                                Some('{') => {
+                                    // Nested `{` inside a field is not allowed.
+                                    return Err(PyValueError::new_err(UNEXPECTED_OPEN));
+                                }
+                                Some('}') => break field_content,
+                                Some(c) => field_content.push(c),
+                            }
+                        };
+
+                        // Split field name from constraint at the first `,`.
+                        let (field_name, new_current_field) =
+                            match closing.find(',') {
+                                Some(comma) => {
+                                    let name = closing[..comma].to_string();
+                                    let full = closing.clone();
+                                    (name, Some(full))
+                                }
+                                None => (closing, None),
+                            };
+
+                        // `_underscore` update based on the last char of the literal.
+                        let new_underscore: Option<String> =
+                            literal.chars().last().map(|last| {
+                                if last == '/' || last == '_' {
+                                    String::new()
+                                } else {
+                                    "_".to_string()
+                                }
+                            });
+
+                        entries.push((
+                            literal,
+                            Some(field_name),
+                            new_underscore,
+                            new_current_field,
+                            true,
+                        ));
+                        literal = String::new();
+                    }
+                }
+            }
+
+            // ---- Closing brace ------------------------------------------
+            Some('}') => {
+                match chars.next() {
+                    Some('}') => {
+                        // `}}` — escaped close brace
+                        literal.push('}');
+                        entries.push((literal, None, Some("_".to_string()), None, false));
+                        literal = String::new();
+                    }
+                    _ => {
+                        // Lone `}` outside a field
+                        return Err(PyValueError::new_err(UNEXPECTED_CLOSE));
+                    }
+                }
+            }
+
+            // ---- Ordinary character -------------------------------------
+            Some(c) => {
+                literal.push(c);
+            }
+        }
+    }
+}
+
+/// Register this module as `snakebids._rust._core`.
+#[pymodule]
+fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(parse_format_string, m)?)?;
+    Ok(())
+}

src/snakebids/_rust/__init__.py

@@ -0,0 +1,5 @@
+"""Private package for optional Rust-accelerated internals.
+
+The compiled extension ``_core`` is built separately with maturin.
+If the extension is not available, the pure-Python fallback is used.
+"""

src/snakebids/utils/snakemake_templates.py

@@ -9,6 +9,13 @@
 import attrs
 from typing_extensions import LiteralString, override
 
+try:
+    from snakebids._rust import _core as _rust_core
+
+    _HAS_RUST_PARSE = True
+except ImportError:
+    _HAS_RUST_PARSE = False
+
 
 @attrs.define(frozen=True)
 class _Wildcard:
@@ -150,11 +157,56 @@ def vformat(
         return super().vformat(format_string, args, kwargs)
 
     @override
-    def parse(  # noqa: PLR0912
+    def parse(
         self, format_string: str
     ) -> Iterator[tuple[str, str | None, str | None, str | None]]:
         """Parse format string, stripping constraints and format specifications.
 
+        When the optional Rust extension (``snakebids._rust._core``) is available the
+        parsing is delegated to a compiled implementation for better performance.
+        Otherwise the pure-Python implementation is used transparently.
+
+        Parameters
+        ----------
+        format_string : str
+            The format string to parse
+
+        Yields
+        ------
+            Each iteration yields (literal_text, field_name, ""|None, None)
+        """
+        if _HAS_RUST_PARSE:
+            yield from self._parse_rust(format_string)
+        else:
+            yield from self._parse_python(format_string)
+
+    def _parse_rust(
+        self, format_string: str
+    ) -> Iterator[tuple[str, str | None, str | None, str | None]]:
+        """Parse using the compiled Rust extension."""
+        try:
+            entries = _rust_core.parse_format_string(format_string)
+        except UnicodeEncodeError:
+            # Python strings may contain lone surrogates (unpaired UTF-16
+            # surrogate code points such as \ud800-\udfff) which are valid
+            # in Python's internal string representation but cannot be
+            # encoded as valid UTF-8 required by Rust's &str.  Fall back to
+            # the pure-Python implementation transparently.
+            yield from self._parse_python(format_string)
+            return
+        for literal, field_name, new_us, new_cf, update_cf in entries:
+            if new_us is not None:
+                self._underscore = new_us
+            if update_cf:
+                self._current_field = new_cf
+            format_spec: str | None = "" if field_name is not None else None
+            yield literal, field_name, format_spec, None
+
+    def _parse_python(  # noqa: PLR0912
+        self, format_string: str
+    ) -> Iterator[tuple[str, str | None, str | None, str | None]]:
+        """Pure-Python parse implementation (fallback when Rust is unavailable).
+
         The function is implemented from scratch in order to avoid the special treatment
         of ``!``, ``:``, and ``[`` by ``string.Formatter.parse()``. It is about 4-5
         times slower than the native implementation, but reasonably well optimized.