Implement binary operations between enums and numbers.

PiperOrigin-RevId: 435352475
Change-Id: I95ffab7d35898b8e04e00e2f37933b7959b90097

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'):

python/mujoco/enums.cc

@@ -12,12 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <cstdint>
 #include <sstream>
 #include <type_traits>
 
 #include "util/crossplatform.h"
 #include "enum_traits.h"
 #include "util/tuple_tools.h"
+#include <pybind11/operators.h>
 #include <pybind11/pybind11.h>
 
 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<int>(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 <typename Tuple>