chore: use factory functions

api/src/opentrons/cli/analyze.py

@@ -30,8 +30,6 @@
 from opentrons_shared_data.robot.dev_types import RobotType
 from opentrons import get_robot_context_tracker
 
-from performance_metrics import RobotContextState
-
 _robot_context_tracker = get_robot_context_tracker()
 
 
@@ -68,7 +66,7 @@ def _get_input_files(files_and_dirs: Sequence[Path]) -> List[Path]:
     return results
 
 
-@_robot_context_tracker.track(state=RobotContextState.ANALYZING_PROTOCOL)
+@_robot_context_tracker.track_analysis()
 async def _analyze(
     files_and_dirs: Sequence[Path],
     json_output: Optional[AsyncPath],

performance-metrics/src/performance_metrics/__init__.py

@@ -1,6 +1,5 @@
 """Opentrons performance metrics library."""
 
-from .datashapes import RobotContextState
 from .robot_context_tracker import RobotContextTracker
 
-__all__ = ["RobotContextState", "RobotContextTracker"]
+__all__ = ["RobotContextTracker"]

performance-metrics/src/performance_metrics/robot_context_tracker.py

@@ -30,7 +30,7 @@ def __init__(self, storage_dir: Path, should_track: bool = False) -> None:
         )
         self._should_track = should_track
 
-    def track(self, state: RobotContextState) -> Callable:  # type: ignore
+    def _track(self, state: RobotContextState) -> Callable:  # type: ignore
         """Decorator factory for tracking the execution duration and state of robot operations.
 
         Args:
@@ -66,6 +66,14 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
 
         return inner_decorator
 
+    def track_analysis(self) -> Callable:  # type: ignore
+        """Decorator for tracking the analysis of a protocol."""
+        return self._track(RobotContextState.ANALYZING_PROTOCOL)
+
+    def track_startup(self) -> Callable:  # type: ignore
+        """Decorator for tracking the startup of the robot."""
+        return self._track(RobotContextState.STARTING_UP)
+
     def store(self) -> None:
         """Returns the stored context data and clears the storage list."""
         stored_data = self._storage.copy()

performance-metrics/tests/performance_metrics/test_robot_context_tracker.py

@@ -31,23 +31,23 @@ def test_file_path(tmp_path: Path, robot_context_tracker: RobotContextTracker) -
 def test_robot_context_tracker(robot_context_tracker: RobotContextTracker) -> None:
     """Tests the tracking of various robot context states through RobotContextTracker."""
 
-    @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
+    @robot_context_tracker._track(state=RobotContextState.STARTING_UP)
     def starting_robot() -> None:
         sleep(STARTING_TIME)
 
-    @robot_context_tracker.track(state=RobotContextState.CALIBRATING)
+    @robot_context_tracker._track(state=RobotContextState.CALIBRATING)
     def calibrating_robot() -> None:
         sleep(CALIBRATING_TIME)
 
-    @robot_context_tracker.track(state=RobotContextState.ANALYZING_PROTOCOL)
+    @robot_context_tracker._track(state=RobotContextState.ANALYZING_PROTOCOL)
     def analyzing_protocol() -> None:
         sleep(ANALYZING_TIME)
 
-    @robot_context_tracker.track(state=RobotContextState.RUNNING_PROTOCOL)
+    @robot_context_tracker._track(state=RobotContextState.RUNNING_PROTOCOL)
     def running_protocol() -> None:
         sleep(RUNNING_TIME)
 
-    @robot_context_tracker.track(state=RobotContextState.SHUTTING_DOWN)
+    @robot_context_tracker._track(state=RobotContextState.SHUTTING_DOWN)
     def shutting_down_robot() -> None:
         sleep(SHUTTING_DOWN_TIME)
 
@@ -85,11 +85,11 @@ def test_multiple_operations_single_state(
 ) -> None:
     """Tests tracking multiple operations within a single robot context state."""
 
-    @robot_context_tracker.track(state=RobotContextState.RUNNING_PROTOCOL)
+    @robot_context_tracker._track(state=RobotContextState.RUNNING_PROTOCOL)
     def first_operation() -> None:
         sleep(RUNNING_TIME)
 
-    @robot_context_tracker.track(state=RobotContextState.RUNNING_PROTOCOL)
+    @robot_context_tracker._track(state=RobotContextState.RUNNING_PROTOCOL)
     def second_operation() -> None:
         sleep(RUNNING_TIME)
 
@@ -111,7 +111,7 @@ def test_exception_handling_in_tracked_function(
 ) -> None:
     """Ensures exceptions in tracked operations are handled correctly."""
 
-    @robot_context_tracker.track(state=RobotContextState.SHUTTING_DOWN)
+    @robot_context_tracker._track(state=RobotContextState.SHUTTING_DOWN)
     def error_prone_operation() -> None:
         sleep(SHUTTING_DOWN_TIME)
         raise RuntimeError("Simulated operation failure")
@@ -133,7 +133,7 @@ async def test_async_operation_tracking(
 ) -> None:
     """Tests tracking of an asynchronous operation."""
 
-    @robot_context_tracker.track(state=RobotContextState.ANALYZING_PROTOCOL)
+    @robot_context_tracker._track(state=RobotContextState.ANALYZING_PROTOCOL)
     async def async_analyzing_operation() -> None:
         await asyncio.sleep(ANALYZING_TIME)
 
@@ -153,7 +153,7 @@ async def test_async_operation_timing_accuracy(
 ) -> None:
     """Tests the timing accuracy of an async operation tracking."""
 
-    @robot_context_tracker.track(state=RobotContextState.RUNNING_PROTOCOL)
+    @robot_context_tracker._track(state=RobotContextState.RUNNING_PROTOCOL)
     async def async_running_operation() -> None:
         await asyncio.sleep(RUNNING_TIME)
 
@@ -172,7 +172,7 @@ async def test_exception_in_async_operation(
 ) -> None:
     """Ensures exceptions in tracked async operations are correctly handled."""
 
-    @robot_context_tracker.track(state=RobotContextState.SHUTTING_DOWN)
+    @robot_context_tracker._track(state=RobotContextState.SHUTTING_DOWN)
     async def async_error_prone_operation() -> None:
         await asyncio.sleep(SHUTTING_DOWN_TIME)
         raise RuntimeError("Simulated async operation failure")
@@ -194,11 +194,11 @@ async def test_concurrent_async_operations(
 ) -> None:
     """Tests tracking of concurrent async operations."""
 
-    @robot_context_tracker.track(state=RobotContextState.CALIBRATING)
+    @robot_context_tracker._track(state=RobotContextState.CALIBRATING)
     async def first_async_calibrating() -> None:
         await asyncio.sleep(CALIBRATING_TIME)
 
-    @robot_context_tracker.track(state=RobotContextState.CALIBRATING)
+    @robot_context_tracker._track(state=RobotContextState.CALIBRATING)
     async def second_async_calibrating() -> None:
         await asyncio.sleep(CALIBRATING_TIME)
 
@@ -217,7 +217,7 @@ def test_no_tracking(tmp_path: Path) -> None:
     """Tests that operations are not tracked when tracking is disabled."""
     robot_context_tracker = RobotContextTracker(tmp_path, should_track=False)
 
-    @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
+    @robot_context_tracker._track(state=RobotContextState.STARTING_UP)
     def operation_without_tracking() -> None:
         sleep(STARTING_TIME)
 
@@ -232,15 +232,15 @@ async def test_storing_to_file(tmp_path: Path) -> None:
     """Tests storing the tracked data to a file."""
     robot_context_tracker = RobotContextTracker(tmp_path, should_track=True)
 
-    @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
+    @robot_context_tracker._track(state=RobotContextState.STARTING_UP)
     def starting_robot() -> None:
         sleep(STARTING_TIME)
 
-    @robot_context_tracker.track(state=RobotContextState.CALIBRATING)
+    @robot_context_tracker._track(state=RobotContextState.CALIBRATING)
     def calibrating_robot() -> None:
         sleep(CALIBRATING_TIME)
 
-    @robot_context_tracker.track(state=RobotContextState.ANALYZING_PROTOCOL)
+    @robot_context_tracker._track(state=RobotContextState.ANALYZING_PROTOCOL)
     def analyzing_protocol() -> None:
         sleep(ANALYZING_TIME)