Implement reversed arithmetic operators for enums and fix division behavior.

PiperOrigin-RevId: 435408435
Change-Id: Ifeeff679cff658d5aae96a9ba5fee19d499656a3

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,

python/mujoco/enums.cc

@@ -16,6 +16,8 @@
 #include <sstream>
 #include <type_traits>
 
+#include <Python.h>
+
 #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 <typename T>
+MUJOCO_ALWAYS_INLINE
+T FloorDiv(T a, T b) {
+  ZeroDenominatorCheck(b);
+  return std::floor(static_cast<double>(a) / static_cast<double>(b));
+}
+
 template <typename Trait>
 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<double>(a) / b;
+        });
+  e.def("__rtruediv__",
+        [](const typename Trait::type& b, double a) -> double {
+          ZeroDenominatorCheck(b);
+          return a / static_cast<double>(b);
+        });
   e.def("__floordiv__",
         [](const typename Trait::type& a, std::int64_t b) -> std::int64_t {
-          return static_cast<int>(a) / b;
+          return FloorDiv<std::int64_t>(a, b);
+        });
+  e.def("__floordiv__",
+        [](const typename Trait::type& a, double b) -> double {
+          return FloorDiv<double>(a, b);
+        });
+  e.def("__rfloordiv__",
+        [](const typename Trait::type& b, std::int64_t a) -> std::int64_t {
+          return FloorDiv<std::int64_t>(a, b);
+        });
+  e.def("__rfloordiv__",
+        [](const typename Trait::type& b, double a) -> double {
+          return FloorDiv<double>(a, b);
+        });
+  e.def("__mod__",
+        [](const typename Trait::type& a, std::int64_t b) -> std::int64_t {
+          return a - FloorDiv<std::int64_t>(a, b) * b;
+        });
+  e.def("__mod__",
+        [](const typename Trait::type& a, double b) -> double {
+          return a - FloorDiv<double>(a, b) * b;
+        });
+  e.def("__rmod__",
+        [](const typename Trait::type& b, std::int64_t a) -> std::int64_t {
+          return a - FloorDiv<std::int64_t>(a, b) * b;
+        });
+  e.def("__rmod__",
+        [](const typename Trait::type& b, double a) -> double {
+          return a - FloorDiv<double>(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());
 }