From df202d23b1eb398ab2585cb240781eecf1f86270 Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Thu, 17 Mar 2022 11:57:18 -0700 Subject: [PATCH] Implement reversed arithmetic operators for enums and fix division behavior. PiperOrigin-RevId: 435408435 Change-Id: Ifeeff679cff658d5aae96a9ba5fee19d499656a3 --- python/mujoco/bindings_test.py | 67 +++++++++++++++++++++++++++++ python/mujoco/enums.cc | 78 +++++++++++++++++++++++++++++++--- 2 files changed, 139 insertions(+), 6 deletions(-) diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 023a96b3..d499d80c 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -663,38 +663,105 @@ Euler integrator, semi-implicit in velocity. with self.assertRaises(ValueError): mujoco.mjtJoint(-1) + def test_enum_as_index(self): + x = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k'] + self.assertEqual(x[mujoco.mjtFrame.mjNFRAME], 'h') + self.assertEqual( + x[mujoco.mjtFrame.mjFRAME_GEOM:mujoco.mjtFrame.mjFRAME_CAMERA], + ['c', 'd']) + 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(7, mujoco.mjtFrame.mjNFRAME) + self.assertEqual(7.0, mujoco.mjtFrame.mjNFRAME) self.assertEqual(mujoco.mjtFrame.mjNFRAME, mujoco.mjtFrame.mjNFRAME) self.assertNotEqual(mujoco.mjtFrame.mjNFRAME, mujoco.mjtFrame.mjFRAME_NONE) + self.assertEqual(-mujoco.mjtFrame.mjNFRAME, -7) + self.assertIsInstance(-mujoco.mjtFrame.mjNFRAME, int) + self.assertEqual(mujoco.mjtFrame.mjNFRAME + 1, 8) self.assertIsInstance(mujoco.mjtFrame.mjNFRAME + 1, int) + self.assertEqual(2 + mujoco.mjtFrame.mjNFRAME, 9) + self.assertIsInstance(2 + mujoco.mjtFrame.mjNFRAME, int) self.assertEqual(mujoco.mjtFrame.mjNFRAME + 1.75, 8.75) + self.assertEqual(2.75 + mujoco.mjtFrame.mjNFRAME, 9.75) self.assertEqual(mujoco.mjtFrame.mjNFRAME - 2, 5) self.assertIsInstance(mujoco.mjtFrame.mjNFRAME - 2, int) + self.assertEqual(8 - mujoco.mjtFrame.mjNFRAME, 1) + self.assertIsInstance(8 - mujoco.mjtFrame.mjNFRAME, int) self.assertEqual(mujoco.mjtFrame.mjNFRAME - 2.25, 4.75) + self.assertEqual(8.25 - mujoco.mjtFrame.mjNFRAME, 1.25) self.assertEqual(mujoco.mjtFrame.mjNFRAME * 3, 21) self.assertIsInstance(mujoco.mjtFrame.mjNFRAME * 3, int) + self.assertEqual(3 * mujoco.mjtFrame.mjNFRAME, 21) + self.assertIsInstance(3 * mujoco.mjtFrame.mjNFRAME, int) self.assertEqual(mujoco.mjtFrame.mjNFRAME * 3.5, 24.5) + self.assertEqual(3.5 * mujoco.mjtFrame.mjNFRAME, 24.5) self.assertEqual(mujoco.mjtFrame.mjNFRAME / 2, 3.5) + self.assertEqual(17.5 / mujoco.mjtFrame.mjNFRAME, 2.5) self.assertEqual(mujoco.mjtFrame.mjNFRAME // 2, 3) self.assertIsInstance(mujoco.mjtFrame.mjNFRAME // 2, int) + self.assertEqual(-mujoco.mjtFrame.mjNFRAME // 2, -4) + self.assertIsInstance(-mujoco.mjtFrame.mjNFRAME // 2, int) + self.assertEqual(20 // mujoco.mjtFrame.mjNFRAME, 2) + self.assertIsInstance(20 // mujoco.mjtFrame.mjNFRAME, int) + self.assertEqual(-20 // mujoco.mjtFrame.mjNFRAME, -3) + self.assertIsInstance(-20 // mujoco.mjtFrame.mjNFRAME, int) + self.assertEqual(mujoco.mjtFrame.mjNFRAME // 2.0, 3) + self.assertIsInstance(mujoco.mjtFrame.mjNFRAME // 2.0, float) + self.assertEqual(-mujoco.mjtFrame.mjNFRAME // 2.0, -4) + self.assertIsInstance(-mujoco.mjtFrame.mjNFRAME // 2.0, float) + self.assertEqual(20.0 // mujoco.mjtFrame.mjNFRAME, 2) + self.assertIsInstance(20.0 // mujoco.mjtFrame.mjNFRAME, float) + self.assertEqual(-20 // mujoco.mjtFrame.mjNFRAME, -3) + self.assertIsInstance(-20.0 // mujoco.mjtFrame.mjNFRAME, float) self.assertEqual(mujoco.mjtFrame.mjNFRAME % 4, 3) self.assertIsInstance(mujoco.mjtFrame.mjNFRAME % 4, int) + self.assertEqual(-mujoco.mjtFrame.mjNFRAME % -4, -3) + self.assertIsInstance(-mujoco.mjtFrame.mjNFRAME % -4, int) + self.assertEqual(-mujoco.mjtFrame.mjNFRAME % 4, 1) + self.assertIsInstance(-mujoco.mjtFrame.mjNFRAME % 4, int) + self.assertEqual(mujoco.mjtFrame.mjNFRAME % -4, -1) + self.assertIsInstance(mujoco.mjtFrame.mjNFRAME % -4, int) + self.assertEqual(9 % mujoco.mjtFrame.mjNFRAME, 2) + self.assertIsInstance(9 % mujoco.mjtFrame.mjNFRAME, int) + self.assertEqual(-9 % -mujoco.mjtFrame.mjNFRAME, -2) + self.assertIsInstance(-9 % -mujoco.mjtFrame.mjNFRAME, int) + self.assertEqual(-9 % mujoco.mjtFrame.mjNFRAME, 5) + self.assertIsInstance(-9 % mujoco.mjtFrame.mjNFRAME, int) + self.assertEqual(9 % -mujoco.mjtFrame.mjNFRAME, -5) + self.assertIsInstance(9 % -mujoco.mjtFrame.mjNFRAME, int) + + with self.assertRaises(ZeroDivisionError): + _ = mujoco.mjtFrame.mjNFRAME / 0 + with self.assertRaises(ZeroDivisionError): + _ = 1 / mujoco.mjtFrame.mjFRAME_NONE + with self.assertRaises(ZeroDivisionError): + _ = mujoco.mjtFrame.mjNFRAME // 0 + with self.assertRaises(ZeroDivisionError): + _ = 1 // mujoco.mjtFrame.mjFRAME_NONE + with self.assertRaises(ZeroDivisionError): + _ = mujoco.mjtFrame.mjNFRAME % 0 + with self.assertRaises(ZeroDivisionError): + _ = 1 % mujoco.mjtFrame.mjFRAME_NONE 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_PASSIVE ^ 33, 1) + self.assertEqual(33 | mujoco.mjtDisableBit.mjDSBL_PASSIVE, 33) + self.assertEqual(33 & mujoco.mjtDisableBit.mjDSBL_PASSIVE, 32) + self.assertEqual(33 ^ mujoco.mjtDisableBit.mjDSBL_PASSIVE, 1) self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL << 1, mujoco.mjtDisableBit.mjDSBL_WARMSTART) self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL >> 3, diff --git a/python/mujoco/enums.cc b/python/mujoco/enums.cc index 3e90bdf6..0c5f6cb9 100644 --- a/python/mujoco/enums.cc +++ b/python/mujoco/enums.cc @@ -16,6 +16,8 @@ #include #include +#include + #include "util/crossplatform.h" #include "enum_traits.h" #include "util/tuple_tools.h" @@ -25,6 +27,22 @@ namespace mujoco::python { namespace { namespace py = ::pybind11; + +MUJOCO_ALWAYS_INLINE +void ZeroDenominatorCheck(double b) { + if (b == 0) { + PyErr_SetString(PyExc_ZeroDivisionError, "division by zero"); + throw py::error_already_set(); + } +} + +template +MUJOCO_ALWAYS_INLINE +T FloorDiv(T a, T b) { + ZeroDenominatorCheck(b); + return std::floor(static_cast(a) / static_cast(b)); +} + template MUJOCO_ALWAYS_INLINE void DefEnum(py::module_& m) { @@ -45,27 +63,75 @@ void DefEnum(py::module_& m) { }), py::arg("value"), py::prepend()); - // Comparison operators - e.def(py::self == double()); - // Arithmetic operators + e.def(-py::self); e.def(py::self + std::int64_t()); e.def(py::self + double()); + e.def(std::int64_t() + py::self); + e.def(double() + py::self); + e.def(py::self - std::int64_t()); e.def(py::self - double()); + e.def(std::int64_t() - py::self); + e.def(double() - py::self); + e.def(py::self * std::int64_t()); e.def(py::self * double()); + e.def(std::int64_t() * py::self); + e.def(double() * py::self); + + e.def("__truediv__", + [](const typename Trait::type& a, double b) -> double { + ZeroDenominatorCheck(b); + return static_cast(a) / b; + }); + e.def("__rtruediv__", + [](const typename Trait::type& b, double a) -> double { + ZeroDenominatorCheck(b); + return a / static_cast(b); + }); e.def("__floordiv__", [](const typename Trait::type& a, std::int64_t b) -> std::int64_t { - return static_cast(a) / b; + return FloorDiv(a, b); + }); + e.def("__floordiv__", + [](const typename Trait::type& a, double b) -> double { + return FloorDiv(a, b); + }); + e.def("__rfloordiv__", + [](const typename Trait::type& b, std::int64_t a) -> std::int64_t { + return FloorDiv(a, b); + }); + e.def("__rfloordiv__", + [](const typename Trait::type& b, double a) -> double { + return FloorDiv(a, b); + }); + e.def("__mod__", + [](const typename Trait::type& a, std::int64_t b) -> std::int64_t { + return a - FloorDiv(a, b) * b; + }); + e.def("__mod__", + [](const typename Trait::type& a, double b) -> double { + return a - FloorDiv(a, b) * b; + }); + e.def("__rmod__", + [](const typename Trait::type& b, std::int64_t a) -> std::int64_t { + return a - FloorDiv(a, b) * b; + }); + e.def("__rmod__", + [](const typename Trait::type& b, double a) -> double { + return a - FloorDiv(a, b) * b; }); - e.def(py::self / double()); - e.def(py::self % std::int64_t()); // Bitwise operators e.def(py::self & std::int64_t()); + e.def(std::int64_t() & py::self); e.def(py::self | std::int64_t()); + e.def(std::int64_t() | py::self); e.def(py::self ^ std::int64_t()); + e.def(std::int64_t() ^ py::self); + + // Bit shifts e.def(py::self << std::int64_t()); e.def(py::self >> std::int64_t()); }