Creating interned pyclass instances

[oscarbenjamin]

I am trying to create a pyclass whose instances would be interned so that creating a new instance with the same internal value would return a pre-existing instance if there is one. I don't see how to do this at all with #[new]. Any tips or other suggestions welcome!

Also is there a mailing list for PyO3?

This Python code demonstrates what I want to achieve:

# pyatom.py

from __future__ import annotations

from weakref import WeakValueDictionary as weakdict


# Intern all Atoms in this weakref dict:
cache: weakdict[int, Atom] = weakdict()


class Atom:

    __slots__ = ("value", "__weakref__")

    value: int

    def __new__(cls, value: int) -> Atom:
        atom = cache.get(value)
        if atom is not None:
            print("Return previous object", atom)
            return atom
        else:
            print("Create new object", atom)
            atom = super().__new__(cls)
            atom.value = value
            # I think setdefault makes this threadsafe
            # (Another thread might insert during this block)
            return cache.setdefault(value, atom)

    def __repr__(self) -> str:
        return f"Atom({self.value})"

With that I can do:

>>> from pyatom import Atom
>>> a1 = Atom(42)
Create new object Atom(42)
>>> a2 = Atom(42)
Return previous object Atom(42)
>>> a1 is a2
True

I'm trying to achieve the same with a pyclass defined in pyo3 but stuck in a few ways. The code I have so far is:

use std::sync::Arc;
use std::sync::RwLock;
use std::collections::HashMap;

use pyo3::prelude::*;

struct Atom {
    value: i32,
}

static CACHE: RwLock::<Option<HashMap<i32, Arc<Atom>>>> = RwLock::new(None);

fn init_cache() {
    let mut cache = CACHE.write().unwrap();
    *cache = Some(HashMap::new());
}

impl Atom {
    fn new(value: i32) -> Arc<Atom> {
        {
            // Access CACHE only for reading
            let cache = CACHE.read().unwrap();
            let atom = cache.as_ref().unwrap().get(&value);
            if let Some(atom) = atom {
                println!("Return old Atom({})", value);
                return atom.clone();
            }
        }
        // Write to CACHE
        println!("Create new Atom({})", value);
        let atom = Arc::new(Atom{ value });
        let mut cache = CACHE.write().unwrap();
        // Race condition: maybe a value already got inserted...
        cache.as_mut().unwrap().insert(value, atom.clone());
        atom
    }
}

#[pyclass(name="Atom")]
struct PyAtom {
    atom: Arc<Atom>,
}

#[pymethods]
impl PyAtom {
    #[new]
    fn py_new(value: i32) -> PyAtom {
        PyAtom{ atom: Atom::new(value) }
    }

    fn __repr__(&self) -> String {
        format!("Atom({})", self.atom.value)
    }
}

#[pymodule]
fn ruatom(_py: Python, m: &PyModule) -> PyResult<()> {
    init_cache();
    m.add_class::<PyAtom>()?;
    Ok(())
}

With that I get:

>>> from ruatom import Atom
>>> a1 = Atom(42)
Create new Atom(42)
>>> a2 = Atom(42)
Return old Atom(42)
>>> a1 is a2
False

I've managed to intern objects at the rust level but the final output False shows that the two created Python objects are distinct. I think I understand why that is but I don't see how I could change PyAtom::py_new to get it to do anything else. The function returns a PyAtom struct which will then be converted by the caller into a Python object and if I understand correctly there is nothing I can return from py_new that would cause Atom(42) to return a pre-existing Python object.

Does anyone have any tips on how to achieve this?

I also realise that my static CACHE variable is not good but I'm not sure what is a better way. Naturally it would seem that I should make it be owned by something like say the type object but I don't see a way to attach it to the type object.

[adamreichold]

Does anyone have any tips on how to achieve this?

I think you need to avoid #[new] for now as PyO3 currently enforces creation of a new Python object. It should be possible to use a #[staticmethod] returning something like Py<Atom> which could be used like a1 = Atom.of(42) from Python.

So maybe something like the following seems to work here:

#[pyclass(frozen)]
struct Atom {
  value: i32,
}

static CACHE: GILOnceCell<Mutex<HashMap<i32, Py<Atom>>>> = GILOnceCell::new();

#[pymethods]
impl Atom {
  #[staticmethod]
  fn of(py: Python<'_>, value: i32) -> Py<Atom> {
     let cache = CACHE.get_or_init(py, || Mutex::new(HashMap::new())).lock().unwrap();

     match cache.entry(value) {
       Entry::Occupied(entry) => entry.into_mut().clone(),
       Entry::Vacant(entry) => {
          let atom = Py::new(py, Atom { value }).unwrap();

          entry.insert(atom).clone()
       }
     }
  }
}

(The Mutex is a bit too much IMHO, as we already operate under the GIL. Starting with PyO3 0.19.0, you should be able to replace it with GILProtected<RefCell<..>> which is more lightweight even though it is a bit of a mouthful.)

Also is there a mailing list for PyO3?

I don't think so. We mostly rely on GitHub discussions and use a Gitter channel. So basically, I think you are at the right place.

I also realise that my static CACHE variable is not good but I'm not sure what is a better way. Naturally it would seem that I should make it be owned by something like say the type object but I don't see a way to attach it to the type object.

You could put it into an attribute of your module using e.g. #[pyo3(pass_module)]. But honestly, we do use static variables internally in PyO3 for now, so let's say you are in good company...

[davidhewitt]

Yes, we have been starting discuss how to avoid using static variables but I think we would need to add some APIs to PyO3 first.

[oscarbenjamin]

So maybe something like the following seems to work here:

Thanks, that was very helpful. From there I got that approach working for Atom.of(42) using the code below.

I think you need to avoid #[new] for now as PyO3 currently enforces creation of a new Python object.

Yes, that makes sense. Is there a way to implement tp_new without using #[new] though?

Yes, we have been starting discuss how to avoid using static variables but I think we would need to add some APIs to PyO3 first.

Does PEP 573 (Module State Access from C Extension Methods) apply here? I think it is difficult to do this in C e.g. here I would want to access this from tp_new but the only relevant object received by tp_new is the PyTypeObject*. So unless the type object somehow has a pointer to the cache I don't know how I can get it without just a global static.

This is how I got it working for now:

use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::sync::Mutex;

use pyo3::prelude::*;
use pyo3::once_cell::GILOnceCell;

#[pyclass(frozen)]
struct Atom {
    value: i32,
}

// Always create new Atoms through the cache.
fn unique_pyatom(py: Python<'_>, value: i32) -> Py<Atom> {
    static CACHE: GILOnceCell<Mutex<HashMap<i32, Py<Atom>>>> = GILOnceCell::new();

    let mut cache = CACHE.get_or_init(py, || Mutex::new(HashMap::new())).lock().unwrap();

    match cache.entry(value) {
        Entry::Occupied(entry) => {
            println!("Return old Atom({})", value);

            entry.into_mut().clone()
        },
        Entry::Vacant(entry) => {
            println!("Create new Atom({})", value);
            let atom = Py::new(py, Atom { value }).unwrap();

            entry.insert(atom).clone()
        }
    }
}

#[pymethods]
impl Atom {
    #[staticmethod]
    fn of(py: Python<'_>, value: i32) -> Py<Atom> {
        unique_pyatom(py, value)
    }

    fn __repr__(&self) -> String {
        format!("Atom({})", self.value)
    }
}

#[pymodule]
fn ruatomv2(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<Atom>()?;
    Ok(())
}

[davidhewitt]

Yes, PEP 573 and that kind of idea. There's a function PyType_GetModuleByDef, which PyO3 might in principle be able to use on newer Python versions.

https://docs.python.org/3.11/c-api/type.html#c.PyType_GetModuleByDef

(Much experimentation needed 😄)