update track_analysis

api/src/opentrons/util/performance_helpers.py

@@ -1,6 +1,7 @@
 """Performance helpers for tracking robot context."""
 
 import inspect
+import functools
 from pathlib import Path
 from opentrons_shared_data.performance.dev_types import (
     SupportsTracking,
@@ -30,19 +31,17 @@ def __init__(self, storage_location: Path, should_track: bool) -> None:
         """Initialize the stubbed tracker."""
         pass
 
-    async def track(
+    def track(
         self,
-        func_to_track: UnderlyingFunction,
-        state: RobotContextState,
-        *args: UnderlyingFunctionParameters.args,
-        **kwargs: UnderlyingFunctionParameters.kwargs
-    ) -> UnderlyingFunctionReturn:
-        """Return the function unchanged."""
+        state: "RobotContextState",
+    ) -> typing.Callable[[UnderlyingFunction], UnderlyingFunction]:
+        """Return the original function."""
+
+        def inner_decorator(func: UnderlyingFunction) -> UnderlyingFunction:
+            """Return the original function."""
+            return func
 
-        if inspect.iscoroutinefunction(func_to_track):
-            return await func_to_track(*args, **kwargs)  # type: ignore
-        else:
-            return func_to_track(*args, **kwargs)  # type: ignore
+        return inner_decorator
 
     def store(self) -> None:
         """Do nothing."""
@@ -76,54 +75,21 @@ def _get_robot_context_tracker() -> SupportsTracking:
     return _robot_context_tracker
 
 
-# def _build_track_wrapper(
-#     func: UnderlyingFunction, state: RobotContextState
-# ) -> WrappedFunction:
-#     """Decorator to track the given state for the decorated function.
-
-#     Args:
-#         func: The function to track.
-#         state: The state of the robot context during the function execution.
-
-#     Returns:
-#         The decorated function.
-#     """
-
-#     async def wrapper(
-#         *args: P.args, **kwargs: P.kwargs
-#     ) -> T:
-#         tracker: SupportsTracking = _get_robot_context_tracker()
-
-#         try:
-
-#             result: T = await tracker.track(
-#                 func_to_track=func, state=state, *args, **kwargs
-#             )
-#         finally:
-#             tracker.store()
-
-#         return result
-
-#     return wrapper
-
-
 def track_analysis(func: UnderlyingFunction) -> UnderlyingFunction:
     """Track the analysis of a protocol and optionally store each run."""
 
-    async def wrapper(
+    @functools.wraps(func)
+    def wrapper(
         *args: UnderlyingFunctionParameters.args,
         **kwargs: UnderlyingFunctionParameters.kwargs
     ) -> UnderlyingFunctionReturn:
         tracker: SupportsTracking = _get_robot_context_tracker()
-
         try:
-
-            result: UnderlyingFunctionReturn = await tracker.track(
-                func_to_track=func, state=RobotContextState.ANALYZING_PROTOCOL
+            result = tracker.track(state=RobotContextState.ANALYZING_PROTOCOL)(func)(
+                *args, **kwargs
             )
         finally:
             tracker.store()
-
         return result
 
     return wrapper