try out new state enum

Author: CamDavidsonPilon

core/pioreactor/background_jobs/base.py

@@ -24,6 +24,7 @@
 from pioreactor.pubsub import Client
 from pioreactor.pubsub import create_client
 from pioreactor.pubsub import QOS
+from pioreactor.state import JobState as states
 from pioreactor.utils import append_signal_handlers
 from pioreactor.utils import get_running_pio_job_id
 from pioreactor.utils import is_pio_job_running
@@ -235,11 +236,11 @@ class _BackgroundJob(metaclass=PostInitCaller):
     """
 
     # Homie lifecycle (normally per device (i.e. an rpi) but we are using it for "nodes", in Homie parlance)
-    INIT: pt.JobState = "init"
-    READY: pt.JobState = "ready"
-    DISCONNECTED: pt.JobState = "disconnected"
-    SLEEPING: pt.JobState = "sleeping"
-    LOST: pt.JobState = "lost"
+    INIT: pt.JobState = states.INIT
+    READY: pt.JobState = states.READY
+    DISCONNECTED: pt.JobState = states.DISCONNECTED
+    SLEEPING: pt.JobState = states.SLEEPING
+    LOST: pt.JobState = states.LOST
 
     # initial state is disconnected, set other metadata
     state = DISCONNECTED
@@ -317,7 +318,7 @@ def __init__(self, unit: pt.Unit, experiment: pt.Experiment, source: str = "app"
         }
 
         # this comes _after_ adding state to published settings
-        self.set_state(self.INIT)
+        self.set_state(states.INIT)
 
         self._set_up_exit_protocol()
         self._blocking_event = threading.Event()
@@ -357,7 +358,7 @@ def __post__init__(self) -> None:
         """
         # setting READY should happen after we write to the job manager, since a job might do a long-running
         # task in on_ready, which delays writing to the db, which means `pio kill` might not see it.
-        self.set_state(self.READY)
+        self.set_state(states.READY)
 
     @property
     def job_key(self):
@@ -494,23 +495,27 @@ def set_state(self, new_state: pt.JobState) -> None:
 
         """
 
-        if new_state not in {self.INIT, self.READY, self.DISCONNECTED, self.SLEEPING, self.LOST}:
+        try:
+            new_state_enum = pt.JobState(new_state)
+        except ValueError:
             self.logger.error(f"saw {new_state}: not a valid state")
             return
 
-        if new_state == self.state:
+        current_state = pt.JobState(self.state)
+
+        if new_state_enum == current_state:
             return
 
-        if hasattr(self, f"on_{self.state}_to_{new_state}"):
+        if hasattr(self, f"on_{current_state}_to_{new_state_enum}"):
             try:
-                getattr(self, f"on_{self.state}_to_{new_state}")()
+                getattr(self, f"on_{current_state}_to_{new_state_enum}")()
             except Exception as e:
-                self.logger.debug(f"Error in on_{self.state}_to_{new_state}")
+                self.logger.debug(f"Error in on_{current_state}_to_{new_state_enum}")
                 self.logger.debug(e, exc_info=True)
                 self.logger.error(e)
                 return
 
-        getattr(self, new_state)()
+        getattr(self, new_state_enum)()
 
     def block_until_disconnected(self) -> None:
         """
@@ -546,8 +551,8 @@ def clean_up(self) -> None:
         """
         Disconnect from brokers, set state to "disconnected", stop any activity.
         """
-        if self.state != self.DISCONNECTED:
-            self.set_state(self.DISCONNECTED)
+        if self.state is not states.DISCONNECTED:
+            self.set_state(states.DISCONNECTED)
         self._clean_up_resources()
 
     def add_to_published_settings(self, setting: str, props: pt.PublishableSetting) -> None:
@@ -740,7 +745,7 @@ def exit_gracefully(reason: int | str, *args) -> None:
                 pass
 
     def init(self) -> None:
-        self.state = self.INIT
+        self.state = states.INIT
 
         try:
             # we delay the specific on_init until after we have done our important protocols.
@@ -754,7 +759,7 @@ def init(self) -> None:
         self._log_state(self.state)
 
     def ready(self) -> None:
-        self.state = self.READY
+        self.state = states.READY
 
         try:
             self.on_ready()
@@ -766,7 +771,7 @@ def ready(self) -> None:
         self._log_state(self.state)
 
     def sleeping(self) -> None:
-        self.state = self.SLEEPING
+        self.state = states.SLEEPING
 
         try:
             self.on_sleeping()
@@ -783,7 +788,7 @@ def lost(self) -> None:
         # 1. Monitor can send a lost signal if `check_against_processes_running` triggers.
         # I think it makes sense to ignore it?
 
-        self.state = self.LOST
+        self.state = states.LOST
         self._log_state(self.state)
 
     def disconnected(self) -> None:
@@ -800,7 +805,7 @@ def disconnected(self) -> None:
             self.logger.debug("Error in on_disconnected:")
             self.logger.debug(e, exc_info=True)
 
-        self.state = self.DISCONNECTED
+        self.state = states.DISCONNECTED
         self._log_state(self.state)
 
         # we "set" the internal event, which will cause any event.waits to end blocking. This should happen last.
@@ -858,7 +863,7 @@ def _publish_defined_settings_to_broker(
                 self._publish_setting(name)
 
     def _log_state(self, state: pt.JobState) -> None:
-        if state in (self.READY, self.DISCONNECTED, self.LOST):
+        if state in {states.READY, states.DISCONNECTED, states.LOST}:
             self.logger.info(state.capitalize() + ".")
         else:
             self.logger.debug(state.capitalize() + ".")
@@ -1155,16 +1160,16 @@ def __post__init__(self):
         )
         super().__post__init__()  # set ready
 
-    def _desired_dodging_mode(self, enable_dodging_od: bool, od_state: str | None) -> bool:
+    def _desired_dodging_mode(self, enable_dodging_od: bool, od_state: pt.JobState | None) -> bool:
         """Return True if we should dodge based on enable flag and OD state."""
         if not enable_dodging_od:
             return False
         # enable_dodging_od is true - user wants it on
         if od_state is None:
             return False
-        if od_state in {self.READY, self.SLEEPING, self.INIT}:
+        if od_state in {states.READY, states.SLEEPING, states.INIT}:
             return True
-        if od_state in {self.LOST, self.DISCONNECTED}:
+        if od_state in {states.LOST, states.DISCONNECTED}:
             return False
         return False
 
@@ -1202,8 +1207,7 @@ def set_currently_dodging_od(self, value: bool):
     def set_enable_dodging_od(self, value: bool):
         """Turn dodging on/off based on user intent, then align mode with current OD state."""
         self.enable_dodging_od = value
-        od_running = is_pio_job_running("od_reading")
-        od_state = self.READY if od_running else self.DISCONNECTED
+        od_state = states.READY if is_pio_job_running("od_reading") else states.DISCONNECTED
 
         desired = self._desired_dodging_mode(self.enable_dodging_od, od_state)
         self.set_currently_dodging_od(desired)
@@ -1213,7 +1217,7 @@ def _od_reading_changed_status(self, state_msg: pt.MQTTMessage) -> None:
         if not self.enable_dodging_od:
             return
 
-        new_state = state_msg.payload.decode()
+        new_state = pt.JobState(state_msg.payload.decode())
         desired = self._desired_dodging_mode(self.enable_dodging_od, new_state)
         self.set_currently_dodging_od(desired)
 

core/pioreactor/calibrations/od_calibration_single_vial.py

@@ -308,7 +308,7 @@ def get_voltage_from_adc() -> pt.Voltage:
                     x_min=minimum_od600,
                     x_max=initial_od600,
                 )
-                st.set_state("sleeping")
+                st.set_state(pt.JobState.SLEEPING)
                 echo()
                 echo(style("Stop❗", fg="red"))
                 echo("Carefully remove vial.")
@@ -320,7 +320,7 @@ def get_voltage_from_adc() -> pt.Voltage:
                 while not confirm(green("Continue?"), default=False):
                     pass
                 current_volume_in_vial = initial_volume_in_vial
-                st.set_state("ready")
+                st.set_state(pt.JobState.READY)
                 st.block_until_rpm_is_close_to_target(abs_tolerance=120)
                 sleep(1.0)
 

core/pioreactor/calibrations/od_calibration_using_standards.py

@@ -218,7 +218,7 @@ def get_voltage_from_adc() -> float:
         od600_values.append(standard_od)
         voltages.append(voltage)
 
-        st.set_state("sleeping")
+        st.set_state(pt.JobState.SLEEPING)
 
         for i in range(len(od600_values)):
             click.clear()
@@ -242,7 +242,7 @@ def get_voltage_from_adc() -> float:
         click.echo()
         while not click.confirm(green("Confirm vial is placed in Pioreactor?"), default=True):
             pass
-        st.set_state("ready")
+        st.set_state(pt.JobState.READY)
         click.echo("Starting stirring.")
         st.block_until_rpm_is_close_to_target(abs_tolerance=120)
         sleep(1.0)

core/pioreactor/state.py

@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+from enum import Enum
+
+
+class JobState(str, Enum):
+    INIT = "init"
+    READY = "ready"
+    SLEEPING = "sleeping"
+    DISCONNECTED = "disconnected"
+    LOST = "lost"
+
+    def __str__(self) -> str:
+        return self.value

core/pioreactor/types.py

@@ -5,6 +5,7 @@
 import typing as t
 
 from msgspec import Meta
+from pioreactor.state import JobState  # noqa: F401
 
 if t.TYPE_CHECKING:
     from pioreactor.pubsub import Client
@@ -96,9 +97,6 @@ class PublishableSetting(t.TypedDict, total=False):
     persist: t.NotRequired[bool]
 
 
-type JobState = t.Literal["init", "ready", "sleeping", "disconnected", "lost"]
-
-
 type LedChannel = t.Literal["A", "B", "C", "D"]
 # these are strings! Don't make them ints, since ints suggest we can perform math on them, that's meaningless.
 # str suggest symbols, which they are.