Background

As part of my ongoing project to reimplement Django’s templating language in Rust, I have been adding support for custom template tags.

Simple tags

The simplest custom tag will look something like:

# time_tags.py
from datetime import datetime
from django import template

register = template.Library()


@register.simple_tag
def time(format_string):
  now = datetime.now()
  return now.strftime(format_string)

This can then be used in a Django template like:

{% load time from time_tags %}
<p>Time: {% time '%H:%M:%S' %}</p>

The context

Django’s templating language uses an object called a context to provide dynamic data to the template renderer. This mostly behaves like a Python dictionary.

details

Technically, Django’s context contains a list of dictionaries. This allows for temporarily changing the value of a variable, for example within a {% for %} loop, while keeping the old value for later use.

A simple tag can be defined that takes the context as the first variable:

# time_tags.py
from datetime import datetime
from django import template

register = template.Library()


@register.simple_tag(
    takes_context=True)
def time(context, format_string):
  timezone = context["timezone"]
  now = datetime.now(tz=timezone)
  return now.strftime(format_string)

Django Rusty Templates

In Django Rusty Templates, I have defined a context as a Rust struct. Here’s a simplified version:

pub struct Context {
  context: HashMap<
    String, Py<PyAny>>,
}

When rendering a template tag, the context is passed to the render method as a mutable reference:

trait Render {
  fn render(
    &self,
    py: Python<'_>,
    context: &mut Context,
  ) -> RenderResult;
}

This is natural when working purely in Rust but custom template tags require passing the context to Python, which doesn’t understand Rust lifetimes.

The standard way of connecting Python and Rust is with PyO3. To pass a Rust type to Python, we can wrap it in a #[pyclass]:

#[pyclass]
struct PyContext {
  context: Context,
}

#[pymethods]
impl PyContext {
  // Methods for Python to read from
  // the context
}

And then we can create this in our render method:

struct CustomTag {
  func: Py<PyAny>,
  takes_context: bool,
}

impl Render for CustomTag {
  fn render(
    &self,
    py: Python<'_>,
    context: &mut Context,
  ) -> RenderResult {
    if self.takes_context {
      let py_context =
        PyContext { context };
      let content = self.func
        .bind(py)
        .call1((py_context,))?;
      Ok(content.to_string())
    }
  }
}

Unfortunately, this doesn’t compile because PyContext requires an owned value, not a mutable reference:

error[E0308]: mismatched types
   --> src/render/tags.rs:807:40
    |
807 | let py_context =
        PyContext { context };
    |               ^^^^^^^
            expected `Context`,
            found `&mut Context`

Turning a mutable reference into an owned value

To make progress, we need to find a way to get an owned version of context. To do this, I turned to std::mem::take, which replaces the data pointed to by &mut context with an empty value (via the Default trait) and returns an owned value:

#[derive(Default)]
pub struct Context {
  context: HashMap<
    String, Py<PyAny>>,
}

impl Render for CustomTag {
  fn render(
    &self,
    py: Python<'_>,
    context: &mut Context,
  ) -> RenderResult {
    if self.takes_context {
      let context =
        std::mem::take(context);
      let py_context =
        PyContext { context };
      let content = self.func
        .call1(py, (py_context,))?;
      Ok(content.to_string())
    }
  }
}

Moving the owned value back into the mutable reference

This works very well for giving Python access to the context. However, once the custom tag’s rendering logic has run we need to regain ownership of the context for use in other Rust tags. To do this, we turn to std::mem::replace:

impl Render for CustomTag {
  fn render(
    &self,
    py: Python<'_>,
    context: &mut Context,
  ) -> RenderResult {
    if self.takes_context {
      let swapped_context =
        std::mem::take(
          swapped_context);
      let py_context = PyContext {
          context: swapped_context
      };
      let content = self.func
        .call1(py, (py_context,))?;
      let _ = std::mem::replace(
        context,
        py_context.context);
      Ok(content.to_string())
    }
  }
}

Unfortunately, this again does not compile:

error[E0382]: use of moved value:
                `py_context.context`
   --> src/render/tags.rs:815:48
    |
813 | let py_context = PyContext { 
        context: swapped_context };
    |     ----------
          move occurs because
          `py_context` has type
          `PyContext`, which does
          not implement the `Copy`
          trait
814 | let content = self.func.call1(
        py, (py_context,))?;
    |        ----------
             value moved here
815 | let _ = std::mem::replace(
        context,
        py_context.context);
    |   ^^^^^^^^^^^^^^^^^^
        value used here after move

To get around this, we can use an Arc (atomic reference count) to send Python a clone of py_context rather than moving it out of scope. We can also remove the context from the Arc with Arc::try_unwrap:

#[pyclass]
#[derive(Clone)]
struct PyContext {
  context: Arc<Context>,
}

impl Render for CustomTag {
  fn render(
    &self,
    py: Python<'_>,
    context: &mut Context,
  ) -> RenderResult {
    if self.takes_context {
      let swapped_context =
        std::mem::take(
          swapped_context).into();
      let py_context = PyContext {
        context: swapped_context
      };
      let content = self.func.call1(
        py, (py_context.clone(),)?;
      let inner_context = match 
          Arc::try_unwrap(
            py_context.context) {
        Ok(inner_context) => {
          inner_context
        }
        Err(_) => todo!(),
      };
      let _ = std::mem::replace(
        context, inner_context);
      Ok(content.to_string())
    }
  }
}

This works great when Python doesn’t keep a reference to the PyContext object. This means the reference count of the Arc is one and Arc::try_unwrap will succeed. If the custom tag implementation keeps a reference around for some reason, we cannot take ownership. Instead we must fall back to cloning the inner context:

#[derive(Default)]
pub struct Context {
  context: HashMap<
    String, Py<PyAny>>,
}

impl Context {
  fn clone_ref(
      &self, py: Python<'_>
  ) -> Self {
    Self {
      context: self
        .context
        .iter()
        .map(|(k, v)| (
          k.clone(),
          v.clone_ref(py),
        ))
        .collect(),
    }
  }
}

impl Render for CustomTag {
  fn render(
    &self,
    py: Python<'_>,
    context: &mut Context,
  ) -> RenderResult {
    if self.takes_context {
      let swapped_context =
        std::mem::take(
          swapped_context).into();
      let py_context = PyContext {
        context: swapped_context
      };
      let content = self.func.call1(
        py, (py_context.clone(),)?;
      let inner_context = match 
          Arc::try_unwrap(
            py_context.context) {
        Ok(inner_context) => {
          inner_context
        }
        Err(inner_context) => {
          inner_context
            .clone_ref(py)
        }
      };
      let _ = std::mem::replace(
        context, inner_context);
      Ok(content.to_string())
    }
  }
}

Note that we need to use the clone_ref method instead of clone because this handles Python’s reference counts correctly.

Mutating the context from Python

This is sufficient to grant Python read-only access to the context, but the context is designed to be mutated. To enable this, we need to protect the context from being mutably accessed from multiple threads. To do this, we can use a Mutex, along with PyO3’s MutexExt trait which provides the lock_py_attached method to avoid deadlocking with the Python interpreter:

use pyo3::sync::MutexExt;

#[pyclass]
#[derive(Clone)]
struct PyContext {
  context: Arc<Mutex<Context>>,
}

impl PyContext {
  fn new(context: Context) -> Self {
    Self {
      context: Arc::new(
        Mutex::new(context)),
    }
  }
}

impl Render for CustomTag {
  fn render(
    &self,
    py: Python<'_>,
    context: &mut Context,
  ) -> RenderResult {
    if self.takes_context {
      let swapped_context =
        std::mem::take(
          swapped_context);
      let py_context =
        PyContext::new(
          swapped_context
      );
      let content = self.func.call1(
        py, (py_context.clone(),)?;
      let inner_context = match 
          Arc::try_unwrap(
            py_context.context) {
        Ok(inner_context) => {
          inner_context
            .into_inner().unwrap()
        }
        Err(inner_context) => {
          let guard = inner_context
            .lock_py_attached(py)
            .unwrap();
          guard.clone_ref(py)
        }
      };
      let _ = std::mem::replace(
        context, inner_context);
      Ok(content.to_string())
    }
  }
}

Conclusions

The inability of PyO3 to expose Rust structs that use lifetimes initially seems limiting, but PyO3 and Rust provide powerful tools to work around these limitations. std::mem::take, std::mem::replace and std::mem::swap allow for advanced manipulation of mutable references and owned values and Arc and Mutex are extremely useful for exposing shared mutable data to Python. PyO3’s MutexExt is essential for working with mutexes and Python together.

You can find the full code implementing a simple custom tag here, with the extra details I omitted here for brevity and clarity.