mypy is fighting back

api/src/opentrons/util/performance_helpers.py

@@ -1,15 +1,16 @@
 """Performance helpers for tracking robot context."""
 
-import asyncio
-import functools
+import inspect
 from pathlib import Path
 from opentrons_shared_data.performance.dev_types import (
  SupportsTracking,
- F,
+ UnderlyingFunction,
+ UnderlyingFunctionParameters,
+ UnderlyingFunctionReturn,
  RobotContextState,
 )
 from opentrons_shared_data.robot.dev_types import RobotTypeEnum
-from typing import Callable, Type
+import typing
 from opentrons.config import (
  get_performance_metrics_data_dir,
  robot_configs,
@@ -21,30 +22,34 @@
  RobotTypeEnum.robot_literal_to_enum(robot_configs.load().model)
 )
 
+
 class StubbedTracker(SupportsTracking):
  """A stubbed tracker that does nothing."""
 
- def __init__(
- self, storage_location: Path, should_track: bool
- ) -> None:
+ def __init__(self, storage_location: Path, should_track: bool) -> None:
  """Initialize the stubbed tracker."""
  pass
 
- def track(self, state: RobotContextState) -> Callable[[F], F]:
+ async def track(
+ self,
+ func_to_track: UnderlyingFunction,
+ state: RobotContextState,
+ *args: UnderlyingFunctionParameters.args,
+ **kwargs: UnderlyingFunctionParameters.kwargs
+ ) -> UnderlyingFunctionReturn:
  """Return the function unchanged."""
 
- def inner_decorator(func: F) -> F:
- """Return the function unchanged."""
- return func
-
- return inner_decorator
+ if inspect.iscoroutinefunction(func_to_track):
+ return await func_to_track(*args, **kwargs) # type: ignore
+ else:
+ return func_to_track(*args, **kwargs) # type: ignore
 
  def store(self) -> None:
  """Do nothing."""
  pass
 
 
-def _handle_package_import() -> Type[SupportsTracking]:
+def _handle_package_import() -> typing.Type[SupportsTracking]:
  """Handle the import of the performance_metrics package.
 
  If the package is not available, return a stubbed tracker.
@@ -71,15 +76,54 @@ def _get_robot_context_tracker() -> SupportsTracking:
  return _robot_context_tracker
 
 
-def track_analysis(func: F) -> F:
+# 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."""
- @functools.wraps(func)
- async def wrapper(*args, **kwargs): # type: ignore # noqa: ANN002, ANN003, ANN201
+
+ async def wrapper(
+ *args: UnderlyingFunctionParameters.args,
+ **kwargs: UnderlyingFunctionParameters.kwargs
+ ) -> UnderlyingFunctionReturn:
  tracker: SupportsTracking = _get_robot_context_tracker()
 
  try:
- return await tracker.track(func, RobotContextState.ANALYZING_PROTOCOL, *args, **kwargs)
+
+ result: UnderlyingFunctionReturn = await tracker.track(
+ func_to_track=func, state=RobotContextState.ANALYZING_PROTOCOL
+ )
  finally:
  tracker.store()
 
- return wrapper # type: ignore
+ return result
+
+ return wrapper

performance-metrics/src/performance_metrics/robot_context_tracker.py

@@ -18,7 +18,7 @@
  SupportsTracking,
  MetricsMetadata,
  UnderlyingFunction,
- UnderlyingFunctionParams,
+ UnderlyingFunctionParameters,
  UnderlyingFunctionReturn,
 )
 
@@ -65,8 +65,8 @@ def __init__(self, storage_location: Path, should_track: bool) -> None:
  async def __call_function(
  self,
  func_to_track: UnderlyingFunction,
- *args: UnderlyingFunctionParams.args,
- **kwargs: UnderlyingFunctionParams.kwargs
+ *args: UnderlyingFunctionParameters.args,
+ **kwargs: UnderlyingFunctionParameters.kwargs
  ) -> UnderlyingFunctionReturn:
  """Calls the given function and returns its result."""
  if inspect.iscoroutinefunction(func_to_track):
@@ -78,8 +78,8 @@ async def track(
  self,
  func_to_track: UnderlyingFunction,
  state: RobotContextState,
- *args: UnderlyingFunctionParams.args,
- **kwargs: UnderlyingFunctionParams.kwargs
+ *args: UnderlyingFunctionParameters.args,
+ **kwargs: UnderlyingFunctionParameters.kwargs
  ) -> UnderlyingFunctionReturn:
  """Tracks the given function and its execution duration.
 

shared-data/python/opentrons_shared_data/performance/dev_types.py

@@ -5,11 +5,9 @@
 from enum import Enum
 
 
+UnderlyingFunction = typing.Callable[..., typing.Awaitable[typing.Any] | typing.Any]
+UnderlyingFunctionParameters = typing.ParamSpec("UnderlyingFunctionParameters")
 UnderlyingFunctionReturn = typing.TypeVar("UnderlyingFunctionReturn")
-UnderlyingFunctionParams = typing.ParamSpec("UnderlyingFunctionParams")
-UnderlyingFunction = typing.TypeVar(
- "UnderlyingFunction", bound=typing.Callable[..., typing.Any]
-)
 
 
 class SupportsTracking(typing.Protocol):
@@ -23,8 +21,8 @@ async def track(
  self,
  func_to_track: UnderlyingFunction,
  state: "RobotContextState",
- *args: UnderlyingFunctionParams.args,
- **kwargs: UnderlyingFunctionParams.kwargs,
+ *args: UnderlyingFunctionParameters.args,
+ **kwargs: UnderlyingFunctionParameters.kwargs,
  ) -> UnderlyingFunctionReturn:
  """Decorator to track the given state for the decorated function."""
  ...