Changed GenericProfiledPIDController to TProfiledPIDController follow Python naming convention. Added Generic annotations and a Protocol to profiledpidcommand.py.

Author: cwstryker

commands2/profiledpidcommand.py

@@ -5,16 +5,17 @@
 # the WPILib BSD license file in the root directory of this project.
 #
 
-from typing import Any, Callable, Union
-
-from .command import Command
-from .subsystem import Subsystem
+from typing import Any, Callable, Generic, Union
 
 from wpimath.controller import ProfiledPIDController, ProfiledPIDControllerRadians
 from wpimath.trajectory import TrapezoidProfile, TrapezoidProfileRadians
 
+from .command import Command
+from .subsystem import Subsystem
+from .typing import TProfiledPIDController, UseOutputFunction
+
 
-class ProfiledPIDCommand(Command):
+class ProfiledPIDCommand(Command, Generic[TProfiledPIDController]):
     """A command that controls an output with a :class:`.ProfiledPIDController`. Runs forever by default -
     to add exit conditions and/or other behavior, subclass this class. The controller calculation and
     output are performed synchronously in the command's execute() method.
@@ -24,10 +25,10 @@ class ProfiledPIDCommand(Command):
 
     def __init__(
         self,
-        controller,
+        controller: TProfiledPIDController,
         measurementSource: Callable[[], float],
         goalSource: Union[float, Callable[[], float]],
-        useOutput: Callable[[float, Any], Any],
+        useOutput: UseOutputFunction,
         *requirements: Subsystem,
     ):
         """Creates a new ProfiledPIDCommand, which controls the given output with a ProfiledPIDController. Goal
@@ -40,14 +41,15 @@ def __init__(
         :param requirements:      the subsystems required by this command
         """
 
+        super().__init__()
         if isinstance(controller, ProfiledPIDController):
             self._stateCls = TrapezoidProfile.State
         elif isinstance(controller, ProfiledPIDControllerRadians):
             self._stateCls = TrapezoidProfileRadians.State
         else:
             raise ValueError(f"unknown controller type {controller!r}")
 
-        self._controller = controller
+        self._controller: TProfiledPIDController = controller
         self._useOutput = useOutput
         self._measurement = measurementSource
         if isinstance(goalSource, (float, int)):

commands2/profiledpidsubsystem.py

@@ -7,10 +7,10 @@
 from wpimath.trajectory import TrapezoidProfile
 
 from .subsystem import Subsystem
-from .typing import GenericProfiledPIDController
+from .typing import TProfiledPIDController
 
 
-class ProfiledPIDSubsystem(Subsystem, Generic[GenericProfiledPIDController]):
+class ProfiledPIDSubsystem(Subsystem, Generic[TProfiledPIDController]):
     """
     A subsystem that uses a :class:`wpimath.controller.ProfiledPIDController`
     or :class:`wpimath.controller.ProfiledPIDControllerRadians` to
@@ -20,12 +20,12 @@ class ProfiledPIDSubsystem(Subsystem, Generic[GenericProfiledPIDController]):
 
     def __init__(
         self,
-        controller: GenericProfiledPIDController,
+        controller: TProfiledPIDController,
         initial_position: float = 0,
     ):
         """Creates a new PIDSubsystem."""
         super().__init__()
-        self._controller: GenericProfiledPIDController = controller
+        self._controller: TProfiledPIDController = controller
         self._enabled = False
         self.setGoal(initial_position)
 
@@ -39,7 +39,7 @@ def periodic(self):
 
     def getController(
         self,
-    ) -> GenericProfiledPIDController:
+    ) -> TProfiledPIDController:
         """Returns the controller"""
         return self._controller
 

commands2/typing.py

@@ -1,7 +1,22 @@
-from typing import TypeVar
+from typing import Protocol, TypeVar
 
 from wpimath.controller import ProfiledPIDController, ProfiledPIDControllerRadians
+from wpimath.trajectory import TrapezoidProfile, TrapezoidProfileRadians
 
-GenericProfiledPIDController = TypeVar(
-    "GenericProfiledPIDController", ProfiledPIDControllerRadians, ProfiledPIDController
+TProfiledPIDController = TypeVar(
+    "TProfiledPIDController", ProfiledPIDControllerRadians, ProfiledPIDController
 )
+TTrapezoidProfileState = TypeVar(
+    "TTrapezoidProfileState",
+    TrapezoidProfileRadians.State,
+    TrapezoidProfile.State,
+)
+
+
+class UseOutputFunction(Protocol):
+
+    def __init__(self): ...
+
+    def __call__(self, t: float, u: TTrapezoidProfileState) -> None: ...
+
+    def accept(self, t: float, u: TTrapezoidProfileState) -> None: ...