diff --git a/src/nb_enum.cpp b/src/nb_enum.cpp index 425008c8..bb261182 100644 --- a/src/nb_enum.cpp +++ b/src/nb_enum.cpp @@ -71,7 +71,7 @@ PyObject *enum_create(enum_init_data *ed) noexcept { scope.attr(name) = result; result.attr("__doc__") = ed->docstr ? str(ed->docstr) : none(); - result.attr("__str__") = enum_mod.attr("Enum").attr("__str__"); + result.attr("__str__") = enum_mod.attr(is_flag ? factory_name : "Enum").attr("__str__"); result.attr("__repr__") = result.attr("__str__"); type_init_data *t = new type_init_data(); @@ -267,6 +267,21 @@ PyObject *enum_from_cpp(const std::type_info *tp, int64_t key) noexcept { return value; } + if ((t->flags & (uint32_t) enum_flags::is_flag) != 0) { + handle enum_tp(t->type_py); + + object val; + if (t->flags & (uint32_t) enum_flags::is_signed) + val = steal(PyLong_FromLongLong((long long) key)); + else + val = steal(PyLong_FromUnsignedLongLong((unsigned long long) key)); + + object result; + result = enum_tp.attr("__new__")(enum_tp, val); + Py_INCREF(result.ptr()); + return result.ptr(); + } + if (t->flags & (uint32_t) enum_flags::is_signed) PyErr_Format(PyExc_ValueError, "%lli is not a valid %s.", (long long) key, t->name); diff --git a/tests/test_enum.cpp b/tests/test_enum.cpp index 5bef7fda..dfafb954 100644 --- a/tests/test_enum.cpp +++ b/tests/test_enum.cpp @@ -49,9 +49,9 @@ NB_MODULE(test_enum_ext, m) { m.def("from_enum", [](Enum value) { return (uint32_t) value; }, nb::arg().noconvert()); m.def("to_enum", [](uint32_t value) { return (Enum) value; }); m.def("from_enum", [](Flag value) { return (uint32_t) value; }, nb::arg().noconvert()); - m.def("to_enum", [](uint32_t value) { return (Flag) value; }); + m.def("to_flag", [](uint32_t value) { return (Flag) value; }); m.def("from_enum", [](SEnum value) { return (int32_t) value; }, nb::arg().noconvert()); - m.def("to_enum", [](uint64_t value) { return (UnsignedFlag) value; }); + m.def("to_unsigned_flag", [](uint64_t value) { return (UnsignedFlag) value; }); m.def("from_enum", [](UnsignedFlag value) { return (uint64_t) value; }, nb::arg().noconvert()); m.def("from_enum_implicit", [](Enum value) { return (uint32_t) value; }); diff --git a/tests/test_enum.py b/tests/test_enum.py index ac75e55f..e39188ce 100644 --- a/tests/test_enum.py +++ b/tests/test_enum.py @@ -170,6 +170,11 @@ def test06_enum_flag(): assert t.from_enum(t.UnsignedFlag.B) == 2 assert t.from_enum(t.UnsignedFlag.All) == 0xffffffffffffffff + assert t.to_flag(1) == t.Flag.A + assert t.to_flag(2) == t.Flag.B + assert t.to_flag(4) == t.Flag.C + assert t.to_flag(5) == (t.Flag.A | t.Flag.C) + def test09_enum_methods(): assert t.Item1.my_value == 0 and t.Item2.my_value == 1 assert t.Item1.get_value() == 0 and t.Item2.get_value() == 1 diff --git a/tests/test_enum_ext.pyi.ref b/tests/test_enum_ext.pyi.ref index dbe89957..37844391 100644 --- a/tests/test_enum_ext.pyi.ref +++ b/tests/test_enum_ext.pyi.ref @@ -100,11 +100,8 @@ def from_enum_implicit(arg: Enum, /) -> int: ... @overload def from_enum_implicit(arg: Flag, /) -> int: ... -@overload def to_enum(arg: int, /) -> Enum: ... -@overload -def to_enum(arg: int, /) -> Flag: ... +def to_flag(arg: int, /) -> Flag: ... -@overload -def to_enum(arg: int, /) -> UnsignedFlag: ... +def to_unsigned_flag(arg: int, /) -> UnsignedFlag: ... diff --git a/tests/test_ndarray_ext.pyi.ref b/tests/test_ndarray_ext.pyi.ref index ad2d87e5..b90384d0 100644 --- a/tests/test_ndarray_ext.pyi.ref +++ b/tests/test_ndarray_ext.pyi.ref @@ -170,8 +170,6 @@ def ret_numpy_const_ref() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2 def ret_numpy_const_ref_f() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2, 4), order='F', writable=False)]: ... -def ret_numpy_half() -> Annotated[ArrayLike, dict(dtype='float16', shape=(2, 4))]: ... - def ret_pytorch() -> Annotated[ArrayLike, dict(dtype='float32', shape=(2, 4))]: ... def ret_tensorflow() -> tensorflow.python.framework.ops.EagerTensor[dtype=float32, shape=(2, 4)]: ...