pybind11/tests/test_cases_for_stubgen.cpp

#include "pybind11/stl.h"
#include "pybind11/stl_bind.h"
#include "pybind11_tests.h"

#include <array>
#include <map>

namespace test_cases_for_stubgen {

// The `basics` code was copied from here (to have all test cases for stubgen in one place):
// https://github.com/python/mypy/blob/c6cb3c6282003dd3dadcf028735f9ba6190a0c84/test-data/pybind11_mypy_demo/src/main.cpp
// Copyright (c) 2016 The Pybind Development Team, All rights reserved.

namespace basics {

int answer() { return 42; }

int sum(int a, int b) { return a + b; }

double midpoint(double left, double right) { return left + (right - left) / 2; }

double weighted_midpoint(double left, double right, double alpha = 0.5) {
    return left + (right - left) * alpha;
}

struct Point {

    enum class LengthUnit { mm = 0, pixel, inch };

    enum class AngleUnit { radian = 0, degree };

    Point() : Point(0, 0) {}
    Point(double x, double y) : x(x), y(y) {}

    static const Point origin;
    static const Point x_axis;
    static const Point y_axis;

    static LengthUnit length_unit;
    static AngleUnit angle_unit;

    double length() const { return std::sqrt(x * x + y * y); }

    double distance_to(double other_x, double other_y) const {
        double dx = x - other_x;
        double dy = y - other_y;
        return std::sqrt(dx * dx + dy * dy);
    }

    double distance_to(const Point &other) const { return distance_to(other.x, other.y); }

    double x, y;
};

const Point Point::origin = Point(0, 0);
const Point Point::x_axis = Point(1, 0);
const Point Point::y_axis = Point(0, 1);

Point::LengthUnit Point::length_unit = Point::LengthUnit::mm;
Point::AngleUnit Point::angle_unit = Point::AngleUnit::radian;

} // namespace basics

void bind_basics(py::module &basics) {

    using namespace basics;

    // Functions
    basics.def(
        "answer", &answer, "answer docstring, with end quote\""); // tests explicit docstrings
    basics.def("sum", &sum, "multiline docstring test, edge case quotes \"\"\"'''");
    basics.def("midpoint", &midpoint, py::arg("left"), py::arg("right"));
    basics.def("weighted_midpoint",
               weighted_midpoint,
               py::arg("left"),
               py::arg("right"),
               py::arg("alpha") = 0.5);

    // Classes
    py::class_<Point> pyPoint(basics, "Point");
    py::enum_<Point::LengthUnit> pyLengthUnit(pyPoint, "LengthUnit");
    py::enum_<Point::AngleUnit> pyAngleUnit(pyPoint, "AngleUnit");

    pyPoint.def(py::init<>())
        .def(py::init<double, double>(), py::arg("x"), py::arg("y"))
        .def("distance_to",
             py::overload_cast<double, double>(&Point::distance_to, py::const_),
             py::arg("x"),
             py::arg("y"))
        .def("distance_to",
             py::overload_cast<const Point &>(&Point::distance_to, py::const_),
             py::arg("other"))
        .def_readwrite("x", &Point::x)
        .def_property(
            "y",
            [](Point &self) { return self.y; },
            [](Point &self, double value) { self.y = value; })
        .def_property_readonly("length", &Point::length)
        .def_property_readonly_static("x_axis", [](py::object /*cls*/) { return Point::x_axis; })
        .def_property_readonly_static("y_axis", [](py::object /*cls*/) { return Point::y_axis; })
        .def_readwrite_static("length_unit", &Point::length_unit)
        .def_property_static(
            "angle_unit",
            [](py::object & /*cls*/) { return Point::angle_unit; },
            [](py::object & /*cls*/, Point::AngleUnit value) { Point::angle_unit = value; });

    pyPoint.attr("origin") = Point::origin;

    pyLengthUnit.value("mm", Point::LengthUnit::mm)
        .value("pixel", Point::LengthUnit::pixel)
        .value("inch", Point::LengthUnit::inch);

    pyAngleUnit.value("radian", Point::AngleUnit::radian)
        .value("degree", Point::AngleUnit::degree);

    // Module-level attributes
    basics.attr("PI") = std::acos(-1);
    basics.attr("__version__") = "0.0.1";
}

struct UserType {
    bool operator<(const UserType &) const { return false; }
};

struct minimal_caster {
    static constexpr auto name = py::detail::const_name<UserType>();

    static py::handle
    cast(UserType const & /*src*/, py::return_value_policy /*policy*/, py::handle /*parent*/) {
        return py::none().release();
    }

    // Maximizing simplicity. This will go terribly wrong for other arg types.
    template <typename>
    using cast_op_type = const UserType &;

    // NOLINTNEXTLINE(google-explicit-constructor)
    operator UserType const &() {
        static UserType obj;
        return obj;
    }

    bool load(py::handle /*src*/, bool /*convert*/) { return false; }
};

} // namespace test_cases_for_stubgen

namespace pybind11 {
namespace detail {

template <>
struct type_caster<test_cases_for_stubgen::UserType> : test_cases_for_stubgen::minimal_caster {};

} // namespace detail
} // namespace pybind11

PYBIND11_MAKE_OPAQUE(std::map<int, test_cases_for_stubgen::UserType>);
PYBIND11_MAKE_OPAQUE(std::map<test_cases_for_stubgen::UserType, int>);
PYBIND11_MAKE_OPAQUE(std::map<float, test_cases_for_stubgen::UserType>);
PYBIND11_MAKE_OPAQUE(std::map<test_cases_for_stubgen::UserType, float>);

TEST_SUBMODULE(cases_for_stubgen, m) {
    auto basics = m.def_submodule("basics");
    test_cases_for_stubgen::bind_basics(basics);

    using UserType = test_cases_for_stubgen::UserType;

    m.def("pass_user_type", [](const UserType &) {});
    m.def("return_user_type", []() { return UserType(); });

    py::bind_map<std::map<int, UserType>>(m, "MapIntUserType");
    py::bind_map<std::map<UserType, int>>(m, "MapUserTypeInt");

#define LOCAL_HELPER(MapTypePythonName, ...)                                                      \
    py::class_<__VA_ARGS__>(m, MapTypePythonName)                                                 \
        .def(                                                                                     \
            "keys",                                                                               \
            [](const __VA_ARGS__ &v) { return py::make_key_iterator(v); },                        \
            py::keep_alive<0, 1>())                                                               \
        .def(                                                                                     \
            "values",                                                                             \
            [](const __VA_ARGS__ &v) { return py::make_value_iterator(v); },                      \
            py::keep_alive<0, 1>())                                                               \
        .def(                                                                                     \
            "__iter__",                                                                           \
            [](const __VA_ARGS__ &v) { return py::make_iterator(v.begin(), v.end()); },           \
            py::keep_alive<0, 1>())

    LOCAL_HELPER("MapFloatUserType", std::map<float, UserType>);
    LOCAL_HELPER("MapUserTypeFloat", std::map<UserType, float>);

#undef MAP_TYPE

    m.def("pass_std_array_int_2", [](const std::array<int, 2> &) {});
    m.def("return_std_array_int_3", []() { return std::array<int, 3>{{1, 2, 3}}; });
}