Support per-argument type signature customization (not just for default value)

[jdumas]

Consider a simple example like this:

m.def(
    "traverse_graph",
    [](const std::vector<std::vector<int>> &graph,
        std::string method) {
        // ...
    },
    "graph"_a,
    "method"_a.sig("typing.Literal['bfs', 'dfs']")
);

Where I would like to override the stub for the "method" argument to use typing.Literal. I would like to avoid overrind the whoe function signature. From what I understand of the current stubgen code, nb::sig() can only be used to customize how the default value is rendered. Could we have a similar annotation to customize how the type itself is rendered? Even without default values.

[wjakob]

You have four options:

• define a subclass of std::string and make a custom type caster for it. Then you have plenty of control over the type signature generation.
• override the function signature using nb::sig("def traverse_graph(...)"). That will replace the whole thing (including defaults), ignoring the function signature generated by nanobind. It does not make sense to combine nb::sig and nb::arg::sig
• define a pattern file. This is a beefed-up structural regexp that is run over the stubs, and which can be used to perform more advanced replacements.
• write your own code that does a pass over the stubs post-build.

The last two options assume that it's mainly the stubs that you want to change. In that case, the __doc__ member would still display the nanobind-generated version.

[wjakob]

If your type replacement/rewriting needs are even more advanced than that, then my gut reaction would be to say it's out of scope for nanobind and better accomplished with a fork.

[jdumas]

Thanks. I think (1) is really the only viable option for me. I want the __doc__ to match the stub, and override the whole function signature is annoying because it gets out of sync easily when updating other function arguments.

Can you explain the rationale for having nb::arg::sig only overriding how the default value is displayed? I'll play around with subclass/type caster in the meantime see if I can achieve what I want.

[jdumas]

I'm struggling a bit to write a derived type caster for std::string. Can you help me get a working example? Here's what I got so far:

#include <nanobind/nanobind.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>

#include <iostream>
#include <string>
#include <vector>

namespace nb = nanobind;

using namespace nb::literals;

struct Literal : public std::string
{
};

NAMESPACE_BEGIN(NB_NAMESPACE)
NAMESPACE_BEGIN(detail)

template <>
struct type_caster<Literal> : public type_caster<std::string>
{
    NB_TYPE_CASTER(Literal, const_name("typing.Literal['bfs', 'dfs']"))
};

NAMESPACE_END(detail)
NAMESPACE_END(NB_NAMESPACE)

void print(const std::string& s)
{
    std::cout << "this: " << s << std::endl;
}

NB_MODULE(nanobind_example, m)
{
    m.doc() = "This is a \"hello world\" example with nanobind";

    m.def(
        "traverse_graph",
        [](Literal method) {
            std::cout << "Traversing graph using method: " << std::endl;
            print(method);
        },
        "method"_a);
}

But when running the code, it just shows an empty string:

Traversing graph using method:
this:

I can get it to work if I copy the implementation of the std::string type caster itself:

template <> struct type_caster<Literal> {
    NB_TYPE_CASTER(Literal, const_name("typing.Literal['bfs', 'dfs']"))

    bool from_python(handle src, uint8_t, cleanup_list *) noexcept {
        Py_ssize_t size;
        const char *str = PyUnicode_AsUTF8AndSize(src.ptr(), &size);
        if (!str) {
            PyErr_Clear();
            return false;
        }
        value = Literal(str, (size_t) size);
        return true;
    }

    static handle from_cpp(const Literal &value, rv_policy,
                           cleanup_list *) noexcept {
        return PyUnicode_FromStringAndSize(value.c_str(), value.size());
    }
};

But I do not understand why simply deriving from type_caster<std::string> doesn't work?

[jdumas]

Looking at the code again, I think deriving from the type_caster<std::string> doesn't work because we end up with two value member variable, one Literal value in the derived class (never set), and one std::string value in the parent class (set by type_caster<std::string>::from_python. So I'm guessing type_caster<std::string> is not meant to be derived from.

[jdumas]

So for reference I managed to write a type-caster that interpret any enum into a typing.Literal[] using magic_enum. Here is a working example:

#include <nanobind/nanobind.h>
#include <nanobind/stl/string.h>

#include <magic_enum/magic_enum.hpp>

#include <array>
#include <iostream>
#include <string>

namespace nb = nanobind;

using namespace nb::literals;

////////////////////////////////////////////////////////////////////////////////

template<class InputIt, class Size, class OutputIt>
constexpr //< since C++20
OutputIt copy_n(InputIt first, Size count, OutputIt result)
{
    if (count > 0)
    {
        *result = *first;
        ++result;
        for (Size i = 1; i != count; ++i, (void)++result)
            *result = *++first;
    }

    return result;
}

////////////////////////////////////////////////////////////////////////////////

template <typename Enum>
struct Literal : public std::string
{
    using std::string::string;

    operator Enum() const
    {
        auto s = magic_enum::enum_cast<Enum>(*this);
        if (!s.has_value()) {
            throw std::runtime_error("Invalid literal value: '" + *this + "'");
        }
        return s.value();
    }

    static constexpr size_t calculate_string_size()
    {
        std::string_view base = "typing.Literal[]";
        size_t total_size = base.size();
        bool first = true;
        for (const auto& name : magic_enum::enum_names<Enum>()) {
            total_size += name.size() + (first ? 2 : 4); // account for quotes and comma+space
            first = false;
        }
        total_size += 1; // null terminator
        return total_size;
    }

    static constexpr auto type_names()
    {
        constexpr size_t N = calculate_string_size();
        char buffer[N];
        std::string_view first = "typing.Literal[";
        std::string_view last = "]";
        copy_n(first.data(), first.size(), buffer);
        size_t offset = first.size();
        for (const auto& name : magic_enum::enum_names<Enum>()) {
            if (offset > first.size()) {
                buffer[offset++] = ',';
                buffer[offset++] = ' ';
            }
            buffer[offset++] = '\'';
            copy_n(name.data(), name.size(), buffer + offset);
            offset += name.size();
            buffer[offset++] = '\'';
        }
        copy_n(last.data(), last.size(), buffer + offset);
        return nb::detail::const_name(buffer);
    }
};

////////////////////////////////////////////////////////////////////////////////

NAMESPACE_BEGIN(NB_NAMESPACE)
NAMESPACE_BEGIN(detail)

template <typename Enum>
struct type_caster<Literal<Enum>>
{
    NB_TYPE_CASTER(Literal<Enum>, Literal<Enum>::type_names())

    bool from_python(handle src, uint8_t, cleanup_list*) noexcept
    {
        Py_ssize_t size;
        const char* str = PyUnicode_AsUTF8AndSize(src.ptr(), &size);
        if (!str) {
            PyErr_Clear();
            return false;
        }
        value = Literal<Enum>(str, (size_t)size);
        return true;
    }

    static handle from_cpp(const Literal<Enum>& value, rv_policy, cleanup_list*) noexcept
    {
        return PyUnicode_FromStringAndSize(value.c_str(), value.size());
    }
};

NAMESPACE_END(detail)
NAMESPACE_END(NB_NAMESPACE)

////////////////////////////////////////////////////////////////////////////////

enum class Method { BFS, DFS };

void print(const std::string& s)
{
    std::cout << "this: " << s << std::endl;
}

void check_enum(const Method& method)
{
    switch (method) {
    case Method::BFS: std::cout << "Method is BFS" << std::endl; break;
    case Method::DFS: std::cout << "Method is DFS" << std::endl; break;
    default: std::cout << "Unknown method" << std::endl; break;
    }
}

NB_MODULE(nanobind_example, m)
{
    m.doc() = "This is a \"hello world\" example with nanobind";

    m.def(
        "traverse_graph",
        [](Literal<Method> method) {
            std::cout << "Traversing graph using method: " << std::endl;
            print(method);
            check_enum(method);
        },
        "method"_a);
}

And the working autocompletion in vscode: