How to tell stubgen about types defined in a Python source files

[nurpax]

Some of the functions in my extension take as input a "mutable bool" (BoolRef). I've declared this type in a .py file that's part of my extension package. (Source of BoolRef at the bottom of this post.)

It's working fine in my extension: I can get and set the value field in my extension's C++ code. The class works naturally in the API I'm binding, and I like that it's defined in a Python file as it's then easy for users of my package to see and understand.

It's all great except that I don't know how to let stubgen know about this type. For example, here's how I'm currently binding a function that uses this BoolRef class:

    m.def("begin", [](const char* name, nb::handle open, ImGuiWindowFlags_ flags) {
        if (!bool_obj.is_none()) {
            bool v = nb::cast<bool>(nb::getattr(bool_obj, "value"));
            bool ret = ImGui::Begin(name, &v, flags);
            nb::setattr(bool_obj, "value", nb::bool_(value));
            return ret;
        } else {
            return ImGui::Begin(name, nullptr, flags);            
        }
    }, "name"_a, "open"_a = nb::none(), "flags"_a.sig("WindowFlags.NONE") = ImGuiWindowFlags_None);

As I declare the open argument to begin() as just nb::handle, stubgen of course has no idea that this is typed BoolRef and will generate the following signature for it:

def begin(name: str, open: object | None = None, flags: WindowFlags = WindowFlags.NONE) -> bool: ...

I can get it to work by rewriting the signature in a pattern file like below, but it'd be nice if could somehow annotate that open is a BoolRef and get rid of my pattern file. Any tips on where to start with this?

slimgui_ext.imgui.begin$:
    \from ...common_types import BoolRef
    def begin(name: str, open: BoolRef | None = None, flags: WindowFlags = WindowFlags.NONE) -> bool: ...

Here's what the BoolRef looks like. This is something I'd prefer to keep as normal Python, as it's easier to modify and easier for users of my package to inspect.

class BaseRef(Generic[T]):
    def __init__(self, value: T):
        self.value = value

    def __repr__(self):
        return f"{self.__class__.__name__}({self.value})"

    def __eq__(self, other):
        return self.value == (other.value if isinstance(other, BaseRef) else other)

    def __ne__(self, other):
        return not self.__eq__(other)

class BoolRef(BaseRef[bool]):
    def __bool__(self):
        return self.value

[wjakob]

This is probably easiest do by making your own wrapper type (nb::object subclass). You can take a look at the nb_types.h header to see how they are implemented. You get to specify how they render in function signatures, and what check is to be done to see if something implements the interface of the wrapper type.

[nurpax]

Thanks!

I got a version of something like this working:

class BoolRef : public nb::object {
    NB_OBJECT_DEFAULT(BoolRef, nb::object, "slimgui.BoolRef", PyType_Check)
};

It's just that I don't know what's the right way to get PyType_Check to pass, as I get the below exception when I try to use these:

File "/Users/janne/dev/slimgui/example/demo_window/demo.py", line 82, in show_demo_window
    if not imgui.begin("Python mini-port of Dear ImGui Demo", open=open_ref, flags=window_flags):
           ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: begin(): incompatible function arguments. The following argument types are supported:
    1. begin(name: str, open: slimgui.BoolRef | None = None, flags: slimgui.slimgui_ext.imgui.WindowFlags = WindowFlags.NONE) -> bool

Invoked with types: str, kwargs = { open: slimgui.BoolRef, flags: slimgui.slimgui_ext.imgui.WindowFlags }

It works if I don't use PyType_Check, f.ex. by replacing it with a function that just returns true. I guess I'll need to zoom into what PyType_Check actually does.

EDIT: heh ok, I guess it makes sense as a BoolRef is not a known type in the same sense as something like a tuple, list might be.

[nurpax]

I figured out a way to check for the correct class instance.

I wonder what's a good approach to caching the results of PyImport_ImportModule and PyObject_GetAttrString to speed up passing BoolRefs to extension method?

static bool is_instance_of_boolref(PyObject *obj) {
    PyObject *mymodule = PyImport_ImportModule("slimgui");
    if (!mymodule) return -1;

    PyObject *myclass = PyObject_GetAttrString(mymodule, "BoolRef");
    Py_DECREF(mymodule);
    if (!myclass) return -1;

    int result = PyObject_IsInstance(obj, myclass);
    Py_DECREF(myclass);
    return result != 0;
}

class BoolRef : public nb::object {
    NB_OBJECT(BoolRef, nb::object, "BoolRef", is_instance_of_boolref)
};