Documenting Native Python Extensions Made With Rust and PyO3

I don’t know how popular it is among software engineers to write native Python extensions in Rust today, but it definitely worth looking at it. PyO3 is a great toolkit for effortlessly developing optimized, GIL-free, memory-safe extensions. I know at least one amazing product written in Rust for Python — Pola.rs, which easily beats Pandas written in pure Python. Likely, we are surrounded by other less visible but shiny pieces of Rust code optimizing Python programs.

I use PyO3 extensively in our software: you may check my Similari, FFmpeg-Rust, or Savant-Rs modules that we use in our computer vision framework, Savant, to find how various problems can be solved with PyO3. However, I’m a toy Rust programmer, so don’t judge the code too critically 😄.

For a long time, I struggled to document the extensions to help users observe only those parts that are available in Python without looking at the Rust source code.

In my first extension, I just wrote a markdown file with API:

Definitely, not the best possible approach. It can work, however, if the exported API is tiny and easy to remember. But recently, I wrote a large extension with many symbols exported to Python — savant-rs. It is a large API; it definitely scratched my head when I thought I needed to tell fellow developers to go to Rust code to read the comments and predict how various symbols will be used from Python. So, I started researching how to document Python-exported pieces of code and generate the documentation for them: unfortunately, there's almost no information about that. However, I stumbled upon a discussion on PyO3 GitHub where a reasonable idea appeared.

Initially, the proposal was to use “.pyi” interface files and document them:

However, I didn’t like the proposed approach because it requires you to do things twice in two different places (Rust code and PYI files, which are optional). Luckily, later in the discussion, someone nicknamed “Haaroon” wrote:

It appeared to be a holy grail of documenting. What is the Napoleon extension? It is an introspection-based documentation extractor developed to automatically document NumPy in Spinx with embedded docstrings.

It was exactly the tool I was looking for, so I went to Haaroon’s profile to find what projects they worked on. There I found Pometry/Raphtory, where they used it. I started experiments with my code and am now ready to share the results with you.

The prologue ends here: I begin to share how to make a well-documented native extension for Python with Rust.

How To

We will use the Sphinx documentation builder to generate beautiful documentation. We need the Napoleon extension to extract the documentation from our native library.

I created a small stub Rust/PyO3 project, where I demonstrated how to tie the pieces of the whole picture together. The project is located on GitHub.

In the project two parts: rust extension library under the extension directory and docs under docs.

The extension is generated cargo new --lib extension , and all PyO3 infrastructure files are filled to support building with Maturin.

The documentation is generated with sphinx-quickstart , but first, the Sphinx requirements must be installed with pip:

cd docs && env pip install -r requirements.txt

Let us begin with the extension documentation and generate HTML for it. I usually like when my extensions include version information right from the crate, so I wrote a simple version() function in extension/src/lib.rs :

#[pyfunction]
fn version() -> String {
    format!("{}", env!("CARGO_PKG_VERSION"))
}

Let us document it. It can be done as follows:

/// Return the version of the Rust crate
///
/// This function is a part of Python module `extension`.
///
/// Returns
/// -------
/// str
///    The version of the Rust crate
///
/// Raises
/// ------
/// None
///   This function does not raise any exceptions
///
/// Example
/// -------
///
/// .. code-block:: python
///
///   from savant_rs.draw_spec import PaddingDraw
///   padding = PaddingDraw(1, 2, 3, 4)
///   padding_copy = padding.copy()
///
#[pyfunction]
fn version() -> String {
    format!("{}", env!("CARGO_PKG_VERSION"))
}

As you can see, we utilize Rust’s documentation blocks but with Sphinx’s RST syntax. This code will not normally display when published to Crates.io.

I want to use the version function in the Sphinx documentation, so I imported it to my docs/source/conf.py as depicted:

Of course, you need to build and install the wheel to make it work. I provide ready-to-use Makefile in the project root:

make

It builds the extension project, installs it in the VENV, and builds the docs. If everything is fine, you can access the documentation with its entry point docs/build/html/index.html. Magic! You must see all the symbols:

Let us document the module itself:

/// This module is a part of Python module `extension`.
///
/// It contains the class :py:class:`MyClass` and the
/// functions :py:func:`add`, :py:func:`version`.
///
/// You can use the regular Sphinx syntax to build relationships between objects.
///
#[pymodule]
fn extension(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<MyClass>()?;
    m.add_function(wrap_pyfunction!(add, m)?)?;
    m.add_function(wrap_pyfunction!(version, m)?)?;
    Ok(())
}

Regenerate the documentation (make), and voila! Now you can see module-level documentation strings:

Let us document the add function:

/// Add two integers
///
/// Arguments
/// ---------
/// a : int
///   The first integer
/// b : int
///   The second integer
///
/// Returns
/// -------
/// int
///   The sum of the two integers
/// 
#[pyfunction]
fn add(a: i32, b: i32) -> i32 {
    a + b
}

Recompile everything with make and update the page:

Now, it is time to document MyClass, begin with the new:

#[pymethods]
impl MyClass {
    /// This is a constructor
    ///
    /// Arguments
    /// ---------
    /// value : int
    ///   The value of the class
    ///
    /// Returns
    /// -------
    /// MyClass
    ///   The class
    ///
    #[new]
    fn new(value: i32) -> Self {
        Self { value }
    }
    ...
}

Regenerate the documentation, aaand… nothing happens. It doesn’t work: you will not see any documentation changes. I wasn’t able to make it work, so my workaround is to have a static method as demonstrated:

#[pymethods]
impl MyClass {
    #[new]
    fn new(value: i32) -> Self {
        Self { value }
    }

    /// This is a constructor
    ///
    /// Arguments
    /// ---------
    /// value : int
    ///   The value of the class
    ///
    /// Returns
    /// -------
    /// MyClass
    ///   The class
    ///
    #[staticmethod]
    fn constructor(value: i32) -> Self {
        Self::new(value)
    }
    ...
}

It works:

You may need a disclaimer for extension users about the workaround. Another way is to document it on the struct level:

/// This is a :py:class:`MyClass` class
///
/// Arguments
/// ---------
/// value : int
///   The value of the class
///
/// Returns
/// -------
/// MyClass
///   The class
///
#[pyclass]
struct MyClass {
    #[pyo3(get)]
    value: i32,
}

It works:

Great, the add method can be documented in the same way as it is done for theadd function so that I will skip it.

Well, we see that we have a struct member value with the getter assigned. How can we document it? Let us begin with a most straightforward approach:

#[pyclass]
struct MyClass {
    /// What are you doing here my little friend?
    ///
    #[pyo3(get,set)]
    value: i32,
}

The result is:

If you implemented getters and setters with #[getter] and #[setter] , annotate only one of them because when both are documented, only the docstring for getter is shown:

#[pyclass]
struct MyClass {
    value: i32,
}

#[pymethods]
impl MyClass {

    /// docstring for get_value
    ///
    #[getter]
    fn get_value(&self) -> PyResult<i32> {
        Ok(self.value)
    }

    /// docstring for set_value **avoid it**!!!
    ///
    #[setter]
    fn set_value(&mut self, value: i32) -> PyResult<()> {
        self.value = value;
        Ok(())
    }
    ...
}

Result:

Basically, you need to explain the use of getter and setter in the getter’s docstring.

Conclusion

Using Sphinx with Napoleon makes it easy to document native Python extensions created with Rust and PyO3. Certain things may not work as expected, and limitations exist, but the tool provides the most convenient and effortless approach to documenting and publishing the source code.

Nb. GitHub’s Copilot is a great help in documenting the code. Thanks to Haaroon for sharing their wisdom.