ENH: Add a __contains__ to AndStatus to tell if it encompasses the sp…

…ecified status
bluesky · Jul 31, 2023 · 1804ba3 · 1804ba3
1 parent 5c413cc
commit 1804ba3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/ophyd/status.py b/ophyd/status.py
@@ -569,6 +569,15 @@ def __str__(self):
             "".format(self.__class__.__name__, self)
         )
 
+    def __contains__(self, status: StatusBase):
+        if self.left == status or self.right == status:
+            return True
+        elif isinstance(self.left, AndStatus):
+            return status in self.left
+        elif isinstance(self.right, AndStatus):
+            return status in self.left
+        return False
+
 
 class Status(StatusBase):
     """

diff --git a/ophyd/tests/test_status.py b/ophyd/tests/test_status.py
@@ -143,6 +143,15 @@ def test_and():
     # make sure deep recursion works
     st4 = st1 & st3
     st5 = st3 & st4
+
+    assert st1 in st3
+    assert st1 in st4
+    assert st1 in st5
+
+    unused_status = StatusBase()
+    assert unused_status not in st3
+    unused_status.set_finished()
+
     state1, cb1 = _setup_state_and_cb()
     state2, cb2 = _setup_state_and_cb()
     state3, cb3 = _setup_state_and_cb()