Skip to content

Commit

Permalink
stubgen: added a __nb_signature__ property for nanobind functions
Browse files Browse the repository at this point in the history 
The ``__nb_signature__`` property of nanobind functions exposes:

- how many overloads a function has

- the function signature for each overload

- the docstring for each overload (or ``None``)

- the default arguments of each overload (if present), which are
  directly returned as Python objects. This leaves the complex and
  potentially ill-defined task of turning those Python objects into
  a valid Python expression to the consumer of this API.

- The returned signature references the default arguments using
  substrings of the form '\0', '\1', etc.
  • Loading branch information
@wjakob
wjakob committed Feb 15, 2024
1 parent bfb41c1 commit d71a990
Show file tree
Hide file tree
Showing 4 changed files with 94 additions and 25 deletions.
99 changes: 78 additions & 21 deletions src/nb_func.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,8 @@ static PyObject *nb_func_vectorcall_simple(PyObject *, PyObject *const *,
size_t, PyObject *) noexcept;
static PyObject *nb_func_vectorcall_complex(PyObject *, PyObject *const *,
size_t, PyObject *) noexcept;
static void nb_func_render_signature(const func_data *f) noexcept;
static uint32_t nb_func_render_signature(const func_data *f,
bool nb_signature_mode = false) noexcept;

int nb_func_traverse(PyObject *self, visitproc visit, void *arg) {
size_t size = (size_t) Py_SIZE(self);
Expand Down Expand Up @@ -860,15 +861,16 @@ PyObject *nb_method_descr_get(PyObject *self, PyObject *inst, PyObject *) {


/// Render the function signature of a single function
static void nb_func_render_signature(const func_data *f) noexcept {
static uint32_t nb_func_render_signature(const func_data *f,
bool nb_signature_mode) noexcept {
const bool is_method = f->flags & (uint32_t) func_flags::is_method,
has_args = f->flags & (uint32_t) func_flags::has_args,
has_var_args = f->flags & (uint32_t) func_flags::has_var_args,
has_var_kwargs = f->flags & (uint32_t) func_flags::has_var_kwargs;

const std::type_info **descr_type = f->descr_types;

uint32_t arg_index = 0;
uint32_t arg_index = 0, n_default_args = 0;
buf.put_dstr(f->name);

for (const char *pc = f->descr; *pc != '\0'; ++pc) {
Expand Down Expand Up @@ -927,26 +929,31 @@ static void nb_func_render_signature(const func_data *f) noexcept {

if (f->args[arg_index].value) {
PyObject *o = f->args[arg_index].value;
PyObject *str = PyObject_Str(o);
bool is_str = PyUnicode_Check(o);

if (str) {
Py_ssize_t size = 0;
const char *cstr =
PyUnicode_AsUTF8AndSize(str, &size);
if (!cstr) {
PyErr_Clear();
if (nb_signature_mode) {
buf.put(" = \\");
buf.put_uint32(n_default_args++);
} else {
PyObject *str = PyObject_Str(o);
bool is_str = PyUnicode_Check(o);

if (str) {
Py_ssize_t size = 0;
const char *cstr =
PyUnicode_AsUTF8AndSize(str, &size);
if (!cstr) {
PyErr_Clear();
} else {
buf.put(" = ");
if (is_str)
buf.put('\'');
buf.put(cstr, (size_t) size);
if (is_str)
buf.put('\'');
}
Py_DECREF(str);
} else {
buf.put(" = ");
if (is_str)
buf.put('\'');
buf.put(cstr, (size_t) size);
if (is_str)
buf.put('\'');
PyErr_Clear();
}
Py_DECREF(str);
} else {
PyErr_Clear();
}
}
}
Expand All @@ -971,9 +978,13 @@ static void nb_func_render_signature(const func_data *f) noexcept {
buf.put('.');
buf.put_dstr((borrow<str>(th.attr("__qualname__"))).c_str());
} else {
if (nb_signature_mode)
buf.put('"');
char *name = type_name(*descr_type);
buf.put_dstr(name);
free(name);
if (nb_signature_mode)
buf.put('"');
}
}

Expand All @@ -989,6 +1000,8 @@ static void nb_func_render_signature(const func_data *f) noexcept {
check(arg_index == f->nargs && !*descr_type,
"nanobind::detail::nb_func_render_signature(%s): arguments inconsistent.",
f->name);

return n_default_args;
}

static PyObject *nb_func_get_name(PyObject *self) {
Expand Down Expand Up @@ -1029,6 +1042,50 @@ static PyObject *nb_func_get_module(PyObject *self) {
}
}

PyObject *nb_func_get_nb_signature(PyObject *self, void *) {
func_data *f = nb_func_data(self);
uint32_t count = (uint32_t) Py_SIZE(self);
PyObject *result = PyTuple_New(count);

for (uint32_t i = 0; i < count; ++i) {
const func_data *fi = f + i;
PyObject *docstr;

if (fi->flags & (uint32_t) func_flags::has_doc && fi->doc[0] != '\0') {
docstr = PyUnicode_FromString(fi->doc);
} else {
docstr = Py_None;
Py_INCREF(docstr);
}

buf.clear();
uint32_t n_default_args = nb_func_render_signature(fi, true),
pos = 2;

PyObject *item = PyTuple_New(2 + n_default_args);
NB_TUPLE_SET_ITEM(item, 0, PyUnicode_FromString(buf.get()));
NB_TUPLE_SET_ITEM(item, 1, docstr);

if (f->flags & (uint32_t) func_flags::has_args) {
for (uint32_t j = 0; j < fi->nargs; ++j) {
PyObject *value = fi->args[j].value;
if (!value)
continue;
Py_INCREF(value);
NB_TUPLE_SET_ITEM(item, pos, value);
pos++;
}
}

check(pos == n_default_args + 2,
"__nb_signature__: default argument counting inconsistency!");

NB_TUPLE_SET_ITEM(result, (Py_ssize_t) i, item);
}

return result;
}

PyObject *nb_func_get_doc(PyObject *self, void *) {
func_data *f = nb_func_data(self);
uint32_t count = (uint32_t) Py_SIZE(self);
Expand Down
2 changes: 2 additions & 0 deletions src/nb_internals.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ NAMESPACE_BEGIN(detail)

extern PyObject *nb_func_getattro(PyObject *, PyObject *);
extern PyObject *nb_func_get_doc(PyObject *, void *);
extern PyObject *nb_func_get_nb_signature(PyObject *, void *);
extern PyObject *nb_bound_method_getattro(PyObject *, PyObject *);
extern int nb_func_traverse(PyObject *, visitproc, void *);
extern int nb_func_clear(PyObject *);
Expand Down Expand Up @@ -117,6 +118,7 @@ static PyMemberDef nb_func_members[] = {

static PyGetSetDef nb_func_getset[] = {
{ "__doc__", nb_func_get_doc, nullptr, nullptr, nullptr },
{ "__nb_signature__", nb_func_get_nb_signature, nullptr, nullptr, nullptr },
{ nullptr, nullptr, nullptr, nullptr, nullptr }
};

Expand Down
8 changes: 4 additions & 4 deletions tests/test_functions.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,10 +143,10 @@ NB_MODULE(test_functions_ext, m) {
});

// Test implicit conversion of various types
m.def("test_11_sl", [](signed long x) { return x; });
m.def("test_11_ul", [](unsigned long x) { return x; });
m.def("test_11_sll", [](signed long long x) { return x; });
m.def("test_11_ull", [](unsigned long long x) { return x; });
m.def("test_11_sl", [](signed long x) { return x; });
m.def("test_11_ul", [](unsigned long x) { return x; });
m.def("test_11_sll", [](signed long long x) { return x; });
m.def("test_11_ull", [](unsigned long long x) { return x; });

// Test string caster
m.def("test_12", [](const char *c) { return nb::str(c); });
Expand Down
10 changes: 10 additions & 0 deletions tests/test_functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -350,3 +350,13 @@ def test39_del():

with pytest.raises(KeyError):
t.test_del_dict({})


def test40_nb_signature():
assert t.test_01.__nb_signature__ == ((r"test_01() -> None", None),)
assert t.test_02.__nb_signature__ == ((r"test_02(j: int = \0, k: int = \1) -> int", None, 8, 1),)
assert t.test_05.__nb_signature__ == ((r"test_05(arg: int, /) -> int", "doc_1"), (r"test_05(arg: float, /) -> int", "doc_2"))
assert t.test_07.__nb_signature__ == (
(r"test_07(arg0: int, arg1: int, /, *args, **kwargs) -> tuple[int, int]", None),
(r"test_07(a: int, b: int, *myargs, **mykwargs) -> tuple[int, int]", None)
)

0 comments on commit d71a990

Please sign in to comment.