diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 9c8668b7..023a96b3 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -663,6 +663,43 @@ Euler integrator, semi-implicit in velocity. with self.assertRaises(ValueError): mujoco.mjtJoint(-1) + def test_enum_ops(self): + # Note: when modifying this test, make sure the enum value is an odd number + # so that the division tests are correctly exercised. + self.assertEqual(mujoco.mjtFrame.mjNFRAME, 7) + self.assertEqual(mujoco.mjtFrame.mjNFRAME, 7.0) + self.assertEqual(mujoco.mjtFrame.mjNFRAME, mujoco.mjtFrame.mjNFRAME) + self.assertNotEqual(mujoco.mjtFrame.mjNFRAME, mujoco.mjtFrame.mjFRAME_NONE) + + self.assertEqual(mujoco.mjtFrame.mjNFRAME + 1, 8) + self.assertIsInstance(mujoco.mjtFrame.mjNFRAME + 1, int) + self.assertEqual(mujoco.mjtFrame.mjNFRAME + 1.75, 8.75) + + self.assertEqual(mujoco.mjtFrame.mjNFRAME - 2, 5) + self.assertIsInstance(mujoco.mjtFrame.mjNFRAME - 2, int) + self.assertEqual(mujoco.mjtFrame.mjNFRAME - 2.25, 4.75) + + self.assertEqual(mujoco.mjtFrame.mjNFRAME * 3, 21) + self.assertIsInstance(mujoco.mjtFrame.mjNFRAME * 3, int) + self.assertEqual(mujoco.mjtFrame.mjNFRAME * 3.5, 24.5) + + self.assertEqual(mujoco.mjtFrame.mjNFRAME / 2, 3.5) + self.assertEqual(mujoco.mjtFrame.mjNFRAME // 2, 3) + self.assertIsInstance(mujoco.mjtFrame.mjNFRAME // 2, int) + + self.assertEqual(mujoco.mjtFrame.mjNFRAME % 4, 3) + self.assertIsInstance(mujoco.mjtFrame.mjNFRAME % 4, int) + + self.assertEqual( + mujoco.mjtDisableBit.mjDSBL_GRAVITY | mujoco.mjtDisableBit.mjDSBL_LIMIT, + 72) + self.assertEqual(mujoco.mjtDisableBit.mjDSBL_PASSIVE | 33, 33) + self.assertEqual(mujoco.mjtDisableBit.mjDSBL_PASSIVE & 33, 32) + self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL << 1, + mujoco.mjtDisableBit.mjDSBL_WARMSTART) + self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL >> 3, + mujoco.mjtDisableBit.mjDSBL_CONTACT) + def test_can_raise_error(self): self.data.pstack = self.data.nstack with self.assertRaisesWithLiteralMatch(mujoco.FatalError, 'Stack overflow'): diff --git a/python/mujoco/enums.cc b/python/mujoco/enums.cc index ced5b4f5..3e90bdf6 100644 --- a/python/mujoco/enums.cc +++ b/python/mujoco/enums.cc @@ -12,12 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include "util/crossplatform.h" #include "enum_traits.h" #include "util/tuple_tools.h" +#include #include namespace mujoco::python { @@ -42,6 +44,30 @@ void DefEnum(py::module_& m) { throw py::value_error(err.str()); }), py::arg("value"), py::prepend()); + + // Comparison operators + e.def(py::self == double()); + + // Arithmetic operators + e.def(py::self + std::int64_t()); + e.def(py::self + double()); + e.def(py::self - std::int64_t()); + e.def(py::self - double()); + e.def(py::self * std::int64_t()); + e.def(py::self * double()); + e.def("__floordiv__", + [](const typename Trait::type& a, std::int64_t b) -> std::int64_t { + return static_cast(a) / b; + }); + e.def(py::self / double()); + e.def(py::self % std::int64_t()); + + // Bitwise operators + e.def(py::self & std::int64_t()); + e.def(py::self | std::int64_t()); + e.def(py::self ^ std::int64_t()); + e.def(py::self << std::int64_t()); + e.def(py::self >> std::int64_t()); } template