Incrementally porting a small Python project to Rust (original) (raw)

I’ve been working on git-branchless, a Git workflow similar to the Mercurial “stacked-diff” workflows used at Google and Facebook. I originally wrote it in Python, but later ported it to Rust. This post details various aspects of the porting process.

Introduction

Motivation

Initially, I prototyped a working version of git-branchless using Python. The git-branchless executable can be invoked in one of two ways:

  1. Explicitly by a user’s command. For example, git-branchless smartlog. This is typically aliased to git sl.
  2. By a Git “hook”. This is an event that triggers when specific actions occur in the repository. For example, when a user makes a commit, it triggers the post-commit hook, which then tells git-branchless to register the commit in its internal database.

However, some Git operations can result in many Git hook invocations, and therefore Python invocations. For example, a rebase of a stack of several commits will trigger several post-commit hooks in sequence. The Python interpreter takes tens or hundreds of milliseconds to start up, which degrades performance in the situation where we invoke it many times serially.

I considered incorporating a long-running background process into the architecture, but I hate the amount of additional complexity and edge-cases associated with such an approach. Instead, I decided to rewrite the project in Rust to address the startup time issue.

Why Rust?

These were the requirements for my choice of language:

I chose Rust because I had some prior experience with it, and I liked the ML-style type system. The Python codebase was already statically-checked with Mypy and written in an ML style, so it was largely a line-by-line port.

Previous Rust experience

Before this project, I had a small amount of experience with Rust. I was working in an OCaml codebase, and implemented one module in Rust for performance reasons. This was also a case of interoperating between Rust and another programming language.

I didn’t have any experience working with a significant amount of Rust code. However, OCaml is similar to Rust (as Rust descends from OCaml), and I had already read a couple of papers on linear types, so the language didn’t surprise me very much.

Why incremental?

The small amount of code in the project could conceivably have been rewritten all at once, but I figured it would take a long time to iron out all the bugs that way. I wanted a stable end result, rather than to have to deal with occasional bugs after the initial porting process, so I preferred to use an incremental approach and port the project one module at a time.

By incremental, I’m referring to the method of porting modules individually, rather than all at once. I accomplished this by calling between Rust and Python as appropriate at runtime.

The last Python-only version of git-branchless is version 0.1.0, and the first Rust-only version is version 0.2.0. You can browse the 65 commits between these versions to see my progress over time.

Strategy

I used the PyO3 library for Rust to handle the Python-Rust interop. See this commit for the initial setup.

PyO3 supports calling Rust from Python and vice-versa:

I began porting the lowest-level modules individually and worked my way up the dependency hierarchy until I arrived at the main function.

Since git-branchless is implemented as a short-lived script, I didn’t face any significant lifetime-related difficulties in the Rust version. The problematic situations arose when the Python and Rust code both needed to keep a reference to a shared resource. I worked around these issues by copying the resources on the Rust side, and then undid those copies when I deleted the Python code. For example, when passing a database connection from Python to Rust, I opened a new database connection to the same database on the Rust side.

IDE ergonomics

If you’re using VS Code, do not use the official Rust extension (rust-lang.rust), as it’s fairly unreliable, and doesn’t offer too many features. Instead use the Rust Analyzer plugin (matklad.rust-analyzer), which is more reliable, and offers many quality-of-life features.

The rust-analyzer IDE experience is refreshingly productive compared to VS Code’s Python offerings. Some feature highlights:

See the manual for a full list.

Build ergonomics

The Rust incremental build times weren’t particularly fast, particularly when using macros from dependencies, as they cannot be compiled ahead of time. Nonetheless, the incremental build time was generally much less than the time it took to run the entire test suite. It was only a problem when selecting only a few tests to run.

On the other hand, the compiler diagnostics, static typing, and IDE support are good enough to write significant amounts of code before having to run a build at all. The typechecker was generally faster than Mypy for Python.

Testing ergonomics

Rust lets you write unit tests inside the same module as the unit itself, which is particularly convenient when the unit is small enough that you don’t want to expose it to callers. Of course, you can write inline tests in Python too — this is more of an improvement versus other static languages like C++ and Java.

The dynamic nature of Python makes testing easier in general. For example, it’s oftentimes useful to mock out a dependency in a test without having to declare an interface for a dependency and inject it. In similar cases in Rust, I created a testing sub-module of the modules in question, which exposed the internals of the module to make them more testable.

The default test runner for Rust leaves a lot to be desired compared to Pytest. It’s difficult to ascertain how many tests ran and how many passed, especially when you’ve attempted to filter out some of the tests. The test runner API is available under a nightly flag, so hopefully we will see some better test runners become available soon.

Interop ergonomics

The ergonomics of the PyO3 library for Python/Rust interop are great. Interop is safe and composable.

Safety: I had no instances of segfaults due to interop. (My only segfault was a stack overflow in Rust-only code, unrelated to interop.) Type mismatches are surfaced as helpful TypeError exceptions in Python. This was a lot better than the OCaml interop I had to do, where not enough type information is carried at runtime to dynamically check type conversions at the interop boundaries.

PyO3 also uses Rust’s ownership system to ensure that you have a handle to Python’s Global Interpreter Lock (GIL) before doing Python operations. Even if you don’t have concurrency in your project, it ensures you have a handle to a correctly-initialized Python runtime before attempting to do Python operations.

Composability: PyO3 makes strong use of Rust’s trait system and coherence rules. If the type system knows how to convert a given set of base types, then it’s easy to teach it how to build an aggregate type consisting of those base types. This allows for lightweight implicit type conversions between Python types and Rust types, without it being overly difficult to figure out how a type is converted.

These conversions are all fallible as well, so you can fail a type conversion when appropriate, which lets you keep boilerplate error-checking out of the calling code.

After having dealt with OCaml’s overly-powerful module system, I can say that I prefer typeclasses (Rust traits) for nearly all day-to-day programming. (The typeclass vs module preference is one reason people prefer Haskell over OCaml or vice-versa.)

PyO3’s hello world

To start, let’s take a look at this complete example:

use pyo3::prelude::*;

fn main() -> PyResult<()> {
    Python::with_gil(|py| {
        let os = py.import("os")?;
        let username: String = os
            .call1("getenv", ("USER",))?
            .extract()?;
        println!("Hello, {}", username);
        Ok(())
    })
}

To summarize:

  1. You create an instance of a Python object using Python::with_gil or similar.
  2. You can then access modules, etc., by calling a method like py.import. Error-checking is handled with the ? operator. It will fail e.g. if the module could not be found.
  3. To call a method with arguments, use the call1 method with a tuple of arguments. The tuple has to contain values of types which implement the IntoPy trait, which is already implemented for Rust built-in types. (The call method works for calls with both positional and keyword arguments; the call0 method works for simple calls without any arguments.) The result is checked with ?. It will fail e.g. if the method threw an exception.
  4. To convert the result into a Rust type, call .extract() on it. The result is checked with ?. It will fail e.g. if the Python type could not be converted into the desired Rust type at runtime. In this case, the desired result type is indicated with the : String type annotation on the let-binding.

It would be difficult to make this simpler and still be statically-checked!

Real example

Consider this example of the Python wrapper for the make_graph function.

#[pyfunction]
fn py_make_graph(
    py: Python,
    repo: PyObject,
    merge_base_db: &PyMergeBaseDb,
    event_replayer: &PyEventReplayer,
    head_oid: Option<PyOidStr>,
    main_branch_oid: PyOid,
    branch_oids: HashSet<PyOidStr>,
    hide_commits: bool,
) -> PyResult<PyCommitGraph> {
    let py_repo = &repo;
    let PyRepo(repo) = repo.extract(py)?;
    let PyMergeBaseDb { merge_base_db } = merge_base_db;
    let PyEventReplayer { event_replayer } = event_replayer;
    let head_oid = HeadOid(head_oid.map(|PyOidStr(oid)| oid));
    let PyOid(main_branch_oid) = main_branch_oid;
    let main_branch_oid = MainBranchOid(main_branch_oid);
    let branch_oids = BranchOids(branch_oids.into_iter().map(|PyOidStr(oid)| oid).collect());
    let graph = make_graph(
        &repo,
        &merge_base_db,
        &event_replayer,
        &head_oid,
        &main_branch_oid,
        &branch_oids,
        hide_commits,
    );
    let graph = map_err_to_py_err(graph, "Could not make graph")?;
    let graph: PyResult<HashMap<PyOidStr, PyNode>> = graph
        .into_iter()
        .map(|(oid, node)| {
            let py_node = node_to_py_node(py, &py_repo, &node)?;
            Ok((PyOidStr(oid), py_node))
        })
        .collect();
    let graph = graph?;
    Ok(graph)
}

Going over it:

  1. The function is annotated with the #[pyfunction] macro. This means that any parameters are automatically converted from PyObjects into the listed types at runtime. It also inserts the py: Python argument, if your function needs it. The corresponding Python function is registered to be available in a module accessible to Python, not shown here.
  2. Most argument types are complex types like PyMergeBaseDb (which wraps the Rust type MergeBaseDb) or composite types like HashSet<PyOidStr>.
    • You can actually annotate regular Rust types to make them convertible into Python types. I chose to create separate types like PyMergeBaseDb so that it would be easier to delete them later, without separating out the Python-specific functionality from the Rust functionality.
    • The type PyOidStr is a wrapper type which converts from a Python str into a Rust git2::Oid. We can’t use a String here, because then PyO3 would convert the Python str into a Rust String, which is not what we want. So we use a wrapper type to define a non-default type conversion behavior.
  3. The repo argument is left as a PyObject, rather than a PyRepo object, because I call Python methods on it later. I get a git2::Repository object corresponding to it in the body of the function with the .extract function.
  4. I explicitly unpack each wrapped argument in the body of the function. This isn’t strictly necessary, and you can just use e.g. repo.0 if you like, which would make the function body significantly shorter. I did this as a stylistic matter to ensure exhaustiveness-checking.
  5. The Rust make_graph function is called with Rust types, and it returns a Rust type.
  6. The result of make_graph is checked with map_err_to_py_err. I’m pretty sure this is unnecessary, and it could have been handled by ?, but I wasn’t familiar enough when I was writing this code.
  7. We convert the result back into a Python type with into_iter().map().collect(), and then check the result with ?.

I could have made this example even shorter by using more wrapper types (for head_oid, main_branch_oid, and branch_oids) and by not creating intermediate variables.

Bugs encountered

I kept track of a selection of bugs encountered while porting from Python to Rust. Nearly all of them were caught by the regular integration tests. The hardest-to-detect ones involved specification mistranslations, such as changing the meaning of domain entities or setting a configuration flag to the wrong default.

One example of a changed domain entity is that the Python git2 library only allows the construction of valid object IDs (OIDs), while the Rust git2 library allows the construction of arbitrary OIDs. This means that invalid OID errors are detected at different times. I could have wrapped the Rust OID type in a wrapper type which forces me to verify that it exists, but I didn’t bother, and instead relied on tests to expose bugs.

These are my raw notes. They do not accurately capture the frequency of each kind of issue.

I filed three issues against Rust libraries that I used:

Results

Porting time

I was generally able to port a given module and its tests in one or two days. However, I worked on this over a long period of time as it’s not part of my day job, so it took a a few months to finish the port (from December 18, 2020 to March 15, 2021).

The Python code was already written in an ML style, with use of algebraic data types and very limited metaprogramming. If you can statically check your code with Mypy, then it will probably be relatively easy to port it to Rust. But if it uses dynamic runtime features, you should expect it to be more difficult.

Lines of code

I used cloc to perform line counts.

Language Type Files Blank Comment Code
Python Source code 16 614 849 2481
Python Test code 14 295 535 1308
Python Total 30 909 1384 3789
Rust Source code 18 531 835 4186
Rust Source noise 18 644
Rust Test code 9 217 11 1642
Rust Test noise 9 277
Rust Total 27 748 846 5828
Rust Total minus noise 27 748 846 4907

The Rust version has significantly more lines of code. A portion of that can be attributed to “line noise” (lines which only consist of whitespace and closing delimiters). I calculated the amount of line noise with this command:

find <dir> -name '*.rs' -exec egrep '^[ });,]+$' {} \; | wc -l

The rest is probably due to more verbose idioms, such as the following:

Despite this, I feel that Rust is nearly as expressive as Python, particularly compared to a language like C++.

Time comparison

It’s hard to compare the speedup between Python and Rust, as I didn’t have any end-to-end benchmarks set up.

I took the following measurements for running the test suites serially. Unfortunately, it’s not very meaningful, because the majority of the time is spent shelling out to Git to set up the test repository.

The specific case of Git hooks taking too long to run was considerably improved. I initiated a rebase of a stack of 20 commits with the Python and Rust versions. The Python version got through only a couple of commits in 10 seconds or so before I cancelled it, whereas the Rust version finished rebasing all the commits in a few seconds. So the particular use-case I was optimizing for was greatly improved.

Conclusions

In summary:

The following are hand-curated posts which you might find interesting.

Want to see more of my posts? Follow me on Twitter or subscribe via RSS.