Having some problems with C++ inheritance and nb::rv_policy::reference_internal

[greglandrum]

Hi,

I'm in the process of porting a fairly large project from boost::python to nanobind and I've hit something I can't figure out. I'm hoping someone here can help.

Before I start on the problem: the port is generally going pretty well. The nanobind docs are really helpful and the library itself is amazingly capable. I anticipated that this was going to be a giant slog, but I'm actually kind of enjoying it. Thanks!

I've distilled the problem I'm having as far as possible to create a reproducible. That's what I include here.

I have three C++ classes that I expose to Python: Atom, QueryAtom, and Mol. QueryAtom is derived from Atom. The Mol class stores a vector of (unique) pointers to Atoms. Here's the code snippet defining those (there's a link to the full file at the bottom of this question)

class Atom {
 public:
  unsigned int atnum = 0;
  Atom() {};
  Atom(unsigned int atnum) : atnum(atnum) {};
  virtual ~Atom() {};
  virtual Atom *copy() const {
    auto res = new Atom(atnum);
    return res;
  }
};

class QueryAtom : public Atom {
 public:
  unsigned int qval = 0;
  QueryAtom() : Atom() {};
  QueryAtom(const Atom &atom) : Atom(atom), qval(atom.atnum) {};
  QueryAtom(const QueryAtom &other) : Atom(other), qval(other.qval) {};
  ~QueryAtom() override {};
  QueryAtom *copy() const override {
    auto res = new QueryAtom(*this);
    return res;
  }
};

class Mol {
 public:
  std::vector<std::unique_ptr<Atom>> atoms;
  Mol() {};
  Mol(const Mol &other) {
    for (const auto &a : other.atoms) {
      atoms.push_back(std::unique_ptr<Atom>(a->copy()));
    }
  }
  void addAtom(unsigned int atnum) {
    atoms.push_back(std::make_unique<Atom>(atnum));
  }
  void addQueryAtom(unsigned int atnum) {
    atoms.push_back(std::make_unique<QueryAtom>(Atom(atnum)));
  }
  Atom *getAtom(unsigned int idx) { return atoms[idx].get(); }
};

These are wrapped for nanobind like this:

NB_MODULE(foo1, m) {
  m.doc() = "Module 1";
  nb::class_<Atom>(m, "Atom", R"DOC(An atom class.)DOC")
      .def(nb::init<>())
      .def(nb::init<unsigned int>(), "idx"_a)
      .def_rw("atnum", &Atom::atnum);
  nb::class_<QueryAtom, Atom>(m, "QueryAtom")
      .def(nb::init<>())
      .def(nb::init<const Atom &>(), "atom"_a)
      .def_rw("qval", &QueryAtom::qval);
  nb::class_<Mol>(m, "Mol", R"DOC(A molecule class.)DOC")
      .def(nb::init<>())
      .def("addAtom", &Mol::addAtom, "atnum"_a)
      .def("addQueryAtom", &Mol::addQueryAtom, "atnum"_a)
      .def("getAtom", &Mol::getAtom, "idx"_a,
           nb::rv_policy::reference_internal);

And then I can use the classes like this and everything works as expected:

from rdkit import foo1

m = foo1.Mol()
m.addAtom(6)
m.addAtom(7)
m.addAtom(8)
m.addQueryAtom(9)

a0 = m.getAtom(0)
a3 = m.getAtom(3)
print(a0.atnum, a3.atnum, a3.qval)

The problem comes when I define a new function that replaces an existing Atom with a QueryAtom and returns a pointer to the replacement as an Atom*.

Here's the C++ function:

Atom *replaceWithQueryAtom(Mol *mol, unsigned int idx) {
  mol->atoms[idx] = std::make_unique<QueryAtom>(*mol->atoms[idx]);
  return mol->atoms[idx].get();
}

I wrap this three different ways (this was part of trying to narrow down the problem):

  m.def(
      "replace",
      [](Mol *mol, unsigned int which) {
        return replaceWithQueryAtom(mol, which);
      },
      "mol"_a, "which"_a, nb::rv_policy::reference_internal);
  m.def(
      "replaceCopy",
      [](Mol *mol, unsigned int which) {
        auto ref = replaceWithQueryAtom(mol, which);
        return ref->copy();
      },
      "mol"_a, "which"_a, nb::rv_policy::take_ownership);
  m.def(
      "replaceLeaky",
      [](Mol *mol, unsigned int which) {
        auto ref = replaceWithQueryAtom(mol, which);
        return ref->copy();
      },
      "mol"_a, "which"_a, nb::rv_policy::reference_internal);

The first of these, replace(), does what I am interested in: replaces the Atom (or QueryAtom) within the molecule with a QueryAtom and returns a pointer to the new QueryAtom as an Atom * using reference_internal.

When I call it, I get the following error:

Traceback (most recent call last):
  File "/scratch/RDKit_git/builds/build_nb/../../Code/RDBoost/nbWrap/foo.py", line 13, in <module>
    qat = foo1.replace(m, 1)  # < generates the error
TypeError: Unable to convert function return value to a Python type! The signature was
    replace(mol: rdkit.foo1.Mol, which: int) -> rdkit.foo1.Atom

The second function, replaceCopy(), which copies the new QueryAtom and returns it as an Atom * with take_ownership, and returns the QueryAtom without problems.

The second function, replaceLeaky(), which copies the new QueryAtom and returns it as an Atom * with reference_internal (yes, I know this will leak... I'm just exploring the boundaries of when the return values work), generates the same exception:

Traceback (most recent call last):
  File "/scratch/RDKit_git/builds/build_nb/../../Code/RDBoost/nbWrap/foo.py", line 15, in <module>
    qat = foo1.replaceLeaky(m, 1)  # < generates the error
TypeError: Unable to convert function return value to a Python type! The signature was
    replaceLeaky(mol: rdkit.foo1.Mol, which: int) -> rdkit.foo1.Atom

Here's sample python (builds on the sample code above) that demonstrates the problem:

# qat = foo1.replace(m, 1)  # < generates the error
qat = foo1.replaceCopy(m, 1)  # < works fine
print(qat.atnum, qat.qval)
qat = foo1.replaceLeaky(m, 1)  # < generates the error

Hopefully I'm doing something basic and stupid so this is easy to fix, but I haven't been able to resolve it.

Here's the whole demo cpp file to make reproducing easy:
foo1.cpp

Here's my test file:
foo.py

[greglandrum]

@wjakob do you have any ideas/suggestions as to what I can do here?

[hpkfft]

It looks like replaceWithQueryAtom() is a free function, not a member function of a class.
So, I'm thinking this is similar to the problematic example.
(Also, I don't see motivation for using lambdas.)

m.def("replace", &replaceWithQueryAtom, "mol"_a, "which"_a, nb::rv_policy::reference)

I did not test the code above. Also, note the C++ side has ownership, but it cannot destroy the object while Python is still using it.

An alternative strategy would be to return a unique pointer or a shared pointer, depending on what you want regarding ownership.

---

It seems you're the author of the C++ code, so another option would be to make replaceWithQueryAtom(unsigned int idx) a public member function of class Mol. (Then, on the C++ side, atoms can be private, which allows implementation changes later if, for example, you want to change std::vector to something else.)
The molecule owns its atoms, and you can use nb::rv_policy::reference_internal.

[greglandrum]

It looks like replaceWithQueryAtom() is a free function, not a member function of a class. So, I'm thinking this is similar to the problematic example. (Also, I don't see motivation for using lambdas.)

m.def("replace", &replaceWithQueryAtom, "mol"_a, "which"_a, nb::rv_policy::reference)

That, combined with nb::keep_alive<1, 0>() seems to have been exactly what I needed to make it work. I need to add the keep_alive becaase in my actual code the atoms maintain a pointer back to their owning molecule, so I don't want the molecule to be cleaned up while that atom still exists.

I missed that reference_internal is only useful for methods. Thanks for pointing that out and demonstrating the solution @hpkfft!