Add color merging methods to Tally

Author: nocarryr

src/tslumd/tallyobj.py

@@ -117,6 +117,62 @@ def from_display(cls, display: 'tslumd.Display', **kwargs) -> 'Tally':
         kw.update({attr:getattr(display, attr) for attr in cls._prop_attrs})
         return cls(display.index, **kw)
 
+    def set_color(self, tally_type: TallyType, color: TallyColor):
+        """Set the color property (or properties) for the given TallyType
+
+        Sets the of :attr:`rh_tally`, :attr:`txt_tally` or :attr:`lh_tally`
+        properties matching the :class:`~.common.TallyType` value
+
+        Arguments:
+            tally_type (TallyType): The :class:`~.common.TallyType` member(s)
+                to set. Multiple types can be specified using
+                bitwise ``|`` operators.
+            color (TallyColor): The :class:`~.common.TallyColor` to set
+
+        .. versionadded:: 0.0.4
+        """
+        for ttype in tally_type:
+            setattr(self, ttype.name, color)
+
+    def merge_color(self, tally_type: TallyType, color: TallyColor):
+        """Merge the color property (or properties) for the given TallyType
+        using the :meth:`set_color` method
+
+        Combines the existing color value with the one provided using a bitwise
+        ``|`` (or) operation
+
+        Arguments:
+            tally_type (TallyType): The :class:`~.common.TallyType` member(s)
+                to merge. Multiple types can be specified using
+                bitwise ``|`` operators.
+            color (TallyColor): The :class:`~.common.TallyColor` to merge
+
+        .. versionadded:: 0.0.4
+        """
+        for ttype in tally_type:
+            cur_color = getattr(self, ttype.name)
+            new_color = cur_color | color
+            if new_color == cur_color:
+                continue
+            self.set_color(ttype, new_color)
+
+    def merge(self, other: 'Tally', tally_type: Optional[TallyType] = TallyType.all_tally):
+        """Merge the color(s) from another Tally instance into this one using
+        the :meth:`merge_color` method
+
+        Arguments:
+            other (Tally): The Tally instance to merge with
+            tally_type (TallyType, optional): The :class:`~.common.TallyType`
+                member(s) to merge. Multiple types can be specified using
+                bitwise ``|`` operators.
+                Default is :attr:`~.common.TallyType.all_tally` (all three types)
+
+        .. versionadded:: 0.0.4
+        """
+        for ttype in tally_type:
+            color = getattr(other, ttype.name)
+            self.merge_color(ttype, color)
+
     def update(self, **kwargs) -> Set[str]:
         """Update any known properties from the given keyword-arguments
 

tests/test_tallyobj.py

@@ -38,6 +38,66 @@ def test_tally_display_conversion(faker):
                 assert disp == Display.from_tally(tally) == tally
                 assert Tally.from_display(disp).normalized_brightness == tally.normalized_brightness
 
+def test_color_merge():
+    t1 = Tally(0)
+
+    t1.set_color(TallyType.rh_tally | TallyType.lh_tally, TallyColor.RED)
+    assert t1.rh_tally == TallyColor.RED
+    assert t1.txt_tally == TallyColor.OFF
+    assert t1.lh_tally == TallyColor.RED
+
+    t1.set_color(TallyType.rh_tally | TallyType.lh_tally, TallyColor.GREEN)
+    assert t1.rh_tally == TallyColor.GREEN
+    assert t1.txt_tally == TallyColor.OFF
+    assert t1.lh_tally == TallyColor.GREEN
+
+    t1.merge_color(TallyType.all_tally, TallyColor.RED)
+    assert t1.rh_tally == TallyColor.AMBER
+    assert t1.txt_tally == TallyColor.RED
+    assert t1.lh_tally == TallyColor.AMBER
+
+    t1.merge_color(TallyType.all_tally, TallyColor.GREEN)
+    assert t1.rh_tally == TallyColor.AMBER
+    assert t1.txt_tally == TallyColor.AMBER
+    assert t1.lh_tally == TallyColor.AMBER
+
+    # Reset t1 to OFF, OFF, RED
+    t1.rh_tally = TallyColor.OFF
+    t1.txt_tally = TallyColor.OFF
+    t1.lh_tally = TallyColor.RED
+
+    # Another Tally with only `txt_tally` set
+    t2 = Tally(1, txt_tally=TallyColor.GREEN)
+
+    # Only `txt_tally` should change
+    t1.merge(t2, TallyType.all_tally)
+    assert t1.rh_tally == TallyColor.OFF
+    assert t1.txt_tally == TallyColor.GREEN
+    assert t1.lh_tally == TallyColor.RED
+
+    # Reset t2 to GREEN, RED, GREEN
+    # t1 is still OFF, GREEN, RED
+    t2.rh_tally = TallyColor.GREEN
+    t2.txt_tally = TallyColor.RED
+    t2.lh_tally = TallyColor.GREEN
+
+    t1.merge(t2, TallyType.rh_tally | TallyType.lh_tally)
+    assert t1.rh_tally == TallyColor.GREEN
+    assert t1.txt_tally == TallyColor.GREEN
+    assert t1.lh_tally == TallyColor.AMBER
+
+    t1.merge(t2, TallyType.all_tally)
+    assert t1.rh_tally == TallyColor.GREEN
+    assert t1.txt_tally == TallyColor.AMBER
+    assert t1.lh_tally == TallyColor.AMBER
+
+    t2.rh_tally = TallyColor.RED
+    t1.merge(t2, TallyType.all_tally)
+    assert t1.rh_tally == TallyColor.AMBER
+    assert t1.txt_tally == TallyColor.AMBER
+    assert t1.lh_tally == TallyColor.AMBER
+
+
 def test_broadcast(faker):
     for _ in range(1000):
         i = faker.pyint(max_value=0xfffe)