xenoform-rs 0.1.2

Inlined rust functions

xenoform-rs

Write and execute superfast rust inside your Python code! Here's how...

Write a type-annotated function or method definition in python, add the rust decorator and put the rust implementation in a docstr:

from xenoform_rs import rust

@rust(py=False)
def vector_sum(v: list[int]) -> int:  # ty: ignore[empty-body]
    """
    Ok(v.iter().sum())
    """

Here's what happens automatically when you import a module with @rust-decorated functions:

First call or after code changes:

1. Rust source code is generated - your Python type signatures are translated to Rust types
2. The extension module is compiled
3. Your decorated Python functions are replaced with their compiled Rust implementations

Subsequent calls: The Rust functions execute directly with minimal overhead.

Change detection: Each module stores a hash of its source code and Cargo.toml. On import, xenoform-rs checks these hashes and automatically rebuilds the module if any changes are detected.

Where files go: By default, the ext subfolder contains binaries, generated source code, and build logs. To change this location see below.

Features

• Supports numpy arrays (via the numpy crate) for customised "vectorised" operations.
• Using annotated types, you can override the default mapping of python types to rust types.
• Supports positional and keyword arguments with defaults, including positional-only and keyword-only markers (/,*)
• Supports *args and **kwargs, mapped (respectively) to &Bound<'py, PyTuple> and Option<&Bound<'py, PyDict>>. NB type annotations for these types are still useful for python type checkers. See test_kwargs.py
• Supports custom dependencies and imports.
• Callable types are supported both as arguments and return values. See below.
• Optional (T | None) types are supported, mapping to Option<T>
• Can link to separate rust sources, see test_modules.py for details.
• By default, supports parallel execution when the python interpreter is free-threaded.

Caveats & points to note:

• callable types (more detail below):
  • only generic (untyped) functions/closures are supported.
  • a type override is necessary to pass functions as arguments. The default works for return values.
• complex: 128 bit support only (i.e. not np.complex64)
• if additional modules are specified, the files are copied into the crate. Modifications to additional modules will trigger a rebuild.
• no support for compound types, other than optional (T | None) (This would require building support for rust enums). Use a type override to a generic python type e.g. Annotated[int | float, "&Bound<'py, PyAny>"] or coerce to a single rust type e.g. Annotated[int | float, "f64"].
• no support currently for linking to external prebuilt binaries
• due to restrictions arising from linguistic differences, xenoform-rs will likely never be as functionally complete than its C++ sister, xenoform
• if you build the modules then change to a different python version, or switch between GIL/freethreaded builds, rust modules will not be importable. A full rebuild is required, so deleting the extension module folder is recommended.

Getting started

Install the package

uv add xenoform-rs  # or pip install xenoform-rs

Simply decorate your rust-implemented functions with the rust decorator factory - it handles all the configuration and compilation. Here's a function that counts the elements in a multidimensional array:

from typing import Annotated

import numpy as np
import numpy.typing as npt

from xenoform_rs import rust, rust_dependency


@rust(
    py=False,  # we don't require the python context as the first argument (we aren't constructing any python objects or calling any python APIs)
    dependencies=[rust_dependency("numpy", version="0.28")],  # declare we need the numpy crate
    imports=["numpy::PyReadonlyArrayDyn"],  # import the type we need
)
def array_nelems(a: npt.NDArray[np.int64]) -> Annotated[int, "usize"]:
    # npt.NDArray[np.int64] maps by default to numpy::PyReadOnlyArrayDyn<i64>
    # the return type is a rust usize which gets converted to a python int
    """
    Ok(a.as_array().shape().iter().product())
    """

if __name__ == "__main__":
    print(array_nelems(np.empty([2, 3, 5, 7], dtype=np.int64)))

The @rust decorator factory parameters

name	type	default	description
py	bool	True	Pass the python context as the first argument. Necessary when (e.g.) creating python objects.
dependencies	list[str] | None	None	Rust package dependencies, the rust_dependency convenience function can be used to specify dependency parameters, e.g. dependencies=[rust_dependency("numpy", version="0.28")].
imports	list[str] | None	None	Additional imports, e.g. imports=["numpy::{PyArray2, PyArrayMethods, PyReadonlyArray2}"]
modules	list[Path | str] | None	None	Sources for additional modules
edition	str	"2024"	The rust edition.
profile	dict[str, str] | None	None	Overrides to (release mode) profile, e.g. optimisation level, strip symbols, etc.
help	str | None	None	Docstring for the function
verbose	bool	False	enable debug logging

Performance

Rust can offer very significant performance enhancements over python, especially where vectorised * operations are not available, but even when they are.

* "vectorisation" in this sense means implementing loops in compiled - rather than interpreted - code. In fact, the compiler also has various optimisations available to it including but by no means limited to "true" vectorisation (meaning hardware SIMD instructions).

The first example deals with an operation on a pandas Series that must be done sequentially, and the second shows that significant performance gains can be had even when a vectorised python implementation is available. Running these examples requires the "examples" optional dependencies (and of course rust):

uv add xenoform-rs --extra examples  # or pip install xenoform-rs[examples]

Loop

This is a Rust vs python comparison of a non-vectorisable sequential operation on a pd.Series. First a python implementation...

def calc_balances_py(data: pd.Series, rate: float) -> pd.Series:
    """Cannot vectorise, since each value is dependent on the previous value"""
    result = pd.Series(index=data.index)
    # Directly access the underlying numpy array for performance. pandas>=3 returns a read only array, so make it writeable
    result_np = result.to_numpy()
    result_np.flags.writeable = True
    current_value = 0.0
    for i, value in data.items():
        current_value = (current_value + value) * (1 - rate)
        result_np[i] = current_value
    return result

...and the equivalent rust implementation. Note that pyo3/rust knows nothing about pandas, but can still work with such objects via their python API:

@rust(
    dependencies=[rust_dependency("numpy", version="0.28")],
    imports=["numpy::{PyArray1, PyArrayMethods}", "pyo3::types::{PyDict, PyAnyMethods}"],
    module_name="loop_rs",  # override as "loop" is a rust keyword
    profile={"strip": "symbols"},
)
def calc_balances_rust(
    data: Annotated[pd.Series, "&Bound<'py, PyAny>"], rate: float
) -> Annotated[pd.Series, "Bound<'py, PyAny>"]:  # ty: ignore[empty-body]
    """

    // extract numpy arrays from the series. Note input is i64, output is f64
    let data_obj = data.call_method0("to_numpy")?;
    let data_np: &Bound<'py, PyArray1<i64>> = data_obj.cast()?;
    let n = data_np.len()?;

    // use the pattern from the numpy documentation
    let result_np = unsafe {
        let r = PyArray1::<f64>::zeros(py, [n], false);
        let mut current_value = 0.0;

        for i in 0..n {
            current_value = (current_value + *data_np.uget([i]) as f64) * (1.0 - rate);
            *r.uget_mut([i]) = current_value;
        }
        r
    };

    // Construct a pd.Series with the same index as the input
    let pd = py.import("pandas")?;
    let kwargs = PyDict::new(py);
    kwargs.set_item("index", data.getattr("index")?)?;
    pd.getattr("Series")?.call((result_np,), Some(&kwargs))

    """

Performance comparison:

N	py (ms)	rust (ms)	speedup
1000	0.5	1.2	-60%
10000	2.0	0.1	2235%
100000	18.7	0.5	3654%
1000000	192.8	2.7	7131%
10000000	1894.8	22.8	8214%

Full code is in examples/loop.py.

Distance Matrix

In this example we compute a distance matrix between $N$ points in $D$ dimensions. An efficient numpy implementation could be:

def calc_dist_matrix_py(p: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    "Compute distance matrix from points, using numpy"
    return np.sqrt(((p[:, np.newaxis, :] - p[np.newaxis, :, :]) ** 2).sum(axis=2))

bearing in mind there is some redundancy here as the resulting matrix is symmetric; however vectorisation with redundancy will always win the tradeoff against loops with no redundancy. But a rust implementation is significantly faster, partly because it can avoid the redundant computations:

@rust(
    dependencies=[rust_dependency("numpy", version="0.28")],
    imports=["numpy::{PyArray2, PyArrayMethods, PyReadonlyArray2}"],
)
def calc_dist_matrix_rust(
    points: Annotated[npt.NDArray[np.float64], "PyReadonlyArray2<f64>"],
) -> Annotated[npt.NDArray[np.float64], "Bound<'py, PyArray2<f64>>"]:
    """

    let points = points.as_array();
    let shape = points.shape();
    let (n, d) = (shape[0], shape[1]);

    let result = PyArray2::zeros(py, [n, n], false);
    let mut r = unsafe { result.as_array_mut() };

    for i in 0..n {
        for j in i + 1..n {
            let mut sum = 0.0;
            for k in 0..d {
                let diff = points.get([i, k]).unwrap() - points.get([j, k]).unwrap();
                sum += diff * diff;
            }
            let dist = sum.sqrt();
            if let Some(x) = r.get_mut([i, j]) {
                *x = dist;
            }
            if let Some(x) = r.get_mut([j, i]) {
                *x = dist;
            }
        }
    }
    Ok(result)

    """

N	py (ms)	rust (ms)	speedup
100	0.4	1.3	-68%
300	3.6	0.2	1907%
1000	28.7	2.3	1162%
3000	208.1	20.8	902%
10000	2270.2	236.2	861%

Full code is in examples/distance_matrix.py.

Type Translations

Default mapping

Basic Python types are recursively mapped to rust types, like so:

Python	rust
None	()
int	i32
np.int32	i32
np.int64	i64
bool	bool
float	f64
np.float32	f32
np.float64	f64
complex	&Bound<'py, PyComplex>
np.complex128	&Bound<'py, PyComplex>
str	String
np.ndarray	PyReadonlyArrayDyn
bytes	&'py [u8]
bytearray	&Bound<'py, PyByteArray
list	Vec
set	HashSet
frozenset	HashSet
dict	HashMap
tuple	(...)
slice	&Bound<'py, PySlice>
Any	&Bound<'py, PyAny>
Self	&Bound<'py, PyAny>
type	&Bound<'py, PyType>
*args	&Bound<'py, PyTuple>
**kwargs	Option<&Bound<'py, PyDict>>
T | None	Option<T>
Callable	&Bound<'py, PyCFunction>
...	&Bound<'py, PyEllipsis>

Thus, dict[str, list[float]] becomes - by default - HashMap<String, Vec<f64>>.

The only type mapped to something mutable is npt.NDArray (PyReadonlyArrayDyn elements are mutable). For dict, list, set or bytearray override to the corresponding pyo3 type, e.g. PyList (see test_inplace.py).

The defaults can be overridden if necessary using the Annotated type, e.g.:

@rust(py=False)
def fibonacci(n: Annotated[int, "u64"]) -> Annotated[int, "u64"]
    ...

Note:

• return types are wrapped in PyResult<> allowing for exceptions to be raised via Err(...). See e.g. test_slice.py
• any &Bound<...> pyo3 type in the return value (even overridden) will have the reference stripped.

Callable Types

Passing and returning functions to and from rust is supported, and they can be used interchangeably with python functions and lambdas.

Annotate types using Callable[...] - this gets mapped to Bound<'py, PyCFunction>. When returning functions, note that pyo3's PyCFunction type does not intrinsically contain information about the function's argument and return types.

For function arguments, the default mapping (to Bound<'py, PyCFunction>) does not support python functions/lambdas. For this reason, use the generic override &Bound<'py, PyAny> (PyAnyMethods implement the call... traits). This example will work with both python and rust functions:

@rust(py=False)
def use_modulo(func: Annotated[Callable[[int], int], "&Bound<'py, PyAny>"], i: int) -> int:
    """
    func.call1((i,))?.extract::<i32>()
    """

See the examples in test_callable.py for more detail.

Configuration

pyo3 version

The pyo3 version can be overridden with the environment variable XENOFORM_RS_PYO3_VERSION. The default - and only supported version - is currently 0.28. Using a different version is not guaranteed to work, and will probably require overrides for all argument and return types.

Location of Extension Modules

By default, compiled modules are placed in an ext subdirectory of your project's root. If this location is unsuitable, it can be overridden using the environment variable XENOFORM_RS_EXTMODULE_ROOT. NB avoid using characters in paths (e.g. space, hyphen) that would not be valid in a python module name.

Free-threaded Interpreter

By default, if the interpreter is free-threaded, extension modules will be built without the GIL. This requires the extension code to be threadsafe. If xenoform detects an environment variable XENOFORM_RS_DISABLE_FT, free-threading is disabled.

Troubleshooting

The generated module source code is written to src/lib.rs in a module-specific folder (e.g. ext/my_module_ext). Cargo build output is redirected to build.log in the that folder. The actual binary will be found in the target/release subfolder.

Adding verbose=True to the rust(...) decorator logs the steps taken, with timings, e.g.:

$ uv run examples/loop.py
08:34:22.535 INFO     registering loop_rs_ext.loop_rs.calc_balances_rust (in ext)
08:34:22.597 INFO     module is up-to-date (d4c7165ade6f52c0aa2ef748c4d6e7c4edce201788a65e7b6e29ebde0d480e3e)
08:34:22.597 INFO     imported compiled module loop_rs
08:34:22.598 INFO     redirected calc_balances_rust to compiled function loop_rs._calc_balances_rust