Change color sensor to return an enum

Use enum.IntEnum to keep backward compatibility with old scripts.

Author: schodet

nxt/sensor/generic.py

@@ -225,6 +225,30 @@ class Color(BaseAnalogSensor):
     light sensor included in the first version or in the Education set.
     """
 
+    class DetectedColor(enum.IntEnum):
+        """Color detected by the sensor.
+
+        This is an :class:`enum.IntEnum` for backward compatibility.
+        """
+
+        BLACK = 1
+        """Black or low intensity."""
+
+        BLUE = 2
+        """Blue."""
+
+        GREEN = 3
+        """Green."""
+
+        YELLOW = 4
+        """Yellow."""
+
+        RED = 5
+        """Red."""
+
+        WHITE = 6
+        """White."""
+
     def __init__(self, brick, port):
         super().__init__(brick, port)
         self.set_light_color(Type.COLOR_FULL)
@@ -258,13 +282,13 @@ def get_color(self):
         """Get detected color.
 
         :return: Detected color.
-        :rtype: int
+        :rtype: DetectedColor
 
         This also sets the sensor mode to color detection (light is cycling quickly to
         determine color).
         """
         self.get_reflected_light(Type.COLOR_FULL)
-        return self.get_input_values().scaled_value
+        return self.DetectedColor(self.get_input_values().scaled_value)
 
     get_sample = get_color
 

tests/test_sensors.py

@@ -129,7 +129,9 @@ def test_color(self, mbrick):
             (Port.S1, True, False, Type.COLOR_RED, Mode.RAW, 114, 46, 46, 46),
             (Port.S1, True, False, Type.COLOR_RED, Mode.RAW, 114, 46, 46, 46),
         ]
-        assert s.get_color() == 4
+        color = s.get_color()
+        assert color == 4
+        assert color == s.DetectedColor.YELLOW
         assert s.get_reflected_light(Type.COLOR_RED) == 46
         assert s.get_light_color() == Type.COLOR_RED
         assert mbrick.mock_calls == [