Make it possible to compare UnitTypeID with UNIT_TYPEID

python-api-src/lib_sc2_typeenums.cpp

@@ -217,7 +217,8 @@ void define_typeenums(py::module & m)
         .value("NEUTRAL_UNBUILDABLEPLATESDESTRUCTIBLE", sc2::UNIT_TYPEID::NEUTRAL_UNBUILDABLEPLATESDESTRUCTIBLE)
         .value("NEUTRAL_UTILITYBOT", sc2::UNIT_TYPEID::NEUTRAL_UTILITYBOT)
         .value("NEUTRAL_VESPENEGEYSER", sc2::UNIT_TYPEID::NEUTRAL_VESPENEGEYSER)
-        .value("NEUTRAL_XELNAGATOWER", sc2::UNIT_TYPEID::NEUTRAL_XELNAGATOWER);
+        .value("NEUTRAL_XELNAGATOWER", sc2::UNIT_TYPEID::NEUTRAL_XELNAGATOWER)
+        .def("__eq__", [](const sc2::UNIT_TYPEID &value, sc2::UnitTypeID &value2) { return value == value2; });
 
     py::enum_<sc2::ABILITY_ID>(m, "ABILITY_ID")
         .value("INVALID", sc2::ABILITY_ID::INVALID)

python-api-src/library.cpp

@@ -50,7 +50,8 @@ PYBIND11_MODULE(library, m)
 
 
     py::class_<sc2::UnitTypeID>(m, "UnitTypeID")
-        .def(py::init<sc2::UNIT_TYPEID>());
+        .def(py::init<sc2::UNIT_TYPEID>())
+        .def("__eq__", [](const sc2::UnitTypeID &value, sc2::UNIT_TYPEID &value2) { return value == value2; });
 
     py::implicitly_convertible<sc2::UNIT_TYPEID, sc2::UnitTypeID>();
 

tests/unittypeid.py

+import unittest
+import sys
+
+sys.path.append('build/python-api-src')
+
+from library import UnitTypeID, UNIT_TYPEID
+
+
+class TestUnitType(unittest.TestCase):
+
+    def test_equality(self):
+        self.assertTrue(UNIT_TYPEID.TERRAN_ARMORY == UNIT_TYPEID.TERRAN_ARMORY)
+    
+    def test_inequality(self):
+        self.assertFalse(UNIT_TYPEID.TERRAN_ARMORY != UNIT_TYPEID.TERRAN_ARMORY)
+
+    def test_convert_equality(self):
+        unit_typeid = UNIT_TYPEID.TERRAN_ARMORY
+        self.assertTrue(unit_typeid == UnitTypeID(unit_typeid))
+        self.assertTrue(UnitTypeID(unit_typeid) == unit_typeid)
+        self.assertFalse(UnitTypeID(unit_typeid) != unit_typeid)
+
+if __name__ == '__main__':
+    unittest.main()
+