diff --git a/commands2/__init__.py b/commands2/__init__.py index 1c243a7..2b64950 100644 --- a/commands2/__init__.py +++ b/commands2/__init__.py @@ -2,6 +2,7 @@ from . import button from . import cmd +from . import typing from .commandscheduler import CommandScheduler from .conditionalcommand import ConditionalCommand diff --git a/commands2/profiledpidcommand.py b/commands2/profiledpidcommand.py index 232eabf..d7ce298 100644 --- a/commands2/profiledpidcommand.py +++ b/commands2/profiledpidcommand.py @@ -5,14 +5,19 @@ # the WPILib BSD license file in the root directory of this project. # -from typing import Any, Callable, Generic, Union +from typing import Any, Generic 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 +from .typing import ( + FloatOrFloatSupplier, + FloatSupplier, + TProfiledPIDController, + UseOutputFunction, +) class ProfiledPIDCommand(Command, Generic[TProfiledPIDController]): @@ -26,8 +31,8 @@ class ProfiledPIDCommand(Command, Generic[TProfiledPIDController]): def __init__( self, controller: TProfiledPIDController, - measurementSource: Callable[[], float], - goalSource: Union[float, Callable[[], float]], + measurementSource: FloatSupplier, + goalSource: FloatOrFloatSupplier, useOutput: UseOutputFunction, *requirements: Subsystem, ): diff --git a/commands2/profiledpidsubsystem.py b/commands2/profiledpidsubsystem.py index 2e970ba..4a05baa 100644 --- a/commands2/profiledpidsubsystem.py +++ b/commands2/profiledpidsubsystem.py @@ -25,7 +25,13 @@ def __init__( controller: TProfiledPIDController, initial_position: float = 0, ): - """Creates a new PIDSubsystem.""" + """ + Creates a new Profiled PID Subsystem using the provided PID Controller + + :param controller: the controller that controls the output + :param initial_position: the initial value of the process variable + + """ super().__init__() self._controller: TProfiledPIDController = controller self._enabled = False @@ -46,15 +52,11 @@ def getController( return self._controller def setGoal(self, goal): - """ - Sets the goal state for the subsystem. - """ + """Sets the goal state for the subsystem.""" self._controller.setGoal(goal) def useOutput(self, output: float, setpoint: TTrapezoidProfileState): - """ - Uses the output from the controller object. - """ + """Uses the output from the controller object.""" raise NotImplementedError(f"{self.__class__} must implement useOutput") def getMeasurement(self) -> float: @@ -75,7 +77,5 @@ def disable(self): self.useOutput(0, TrapezoidProfile.State()) def isEnabled(self) -> bool: - """ - Returns whether the controller is enabled. - """ + """Returns whether the controller is enabled.""" return self._enabled diff --git a/commands2/typing.py b/commands2/typing.py index 357b928..6168330 100644 --- a/commands2/typing.py +++ b/commands2/typing.py @@ -1,8 +1,10 @@ -from typing import Protocol, TypeVar +from typing import Callable, Protocol, TypeVar, Union +from typing_extensions import TypeAlias from wpimath.controller import ProfiledPIDController, ProfiledPIDControllerRadians from wpimath.trajectory import TrapezoidProfile, TrapezoidProfileRadians +# Generic Types TProfiledPIDController = TypeVar( "TProfiledPIDController", ProfiledPIDControllerRadians, ProfiledPIDController ) @@ -13,10 +15,16 @@ ) +# Protocols - Structural Typing class UseOutputFunction(Protocol): - def __init__(self): ... + def __init__(self, *args, **kwargs) -> None: ... def __call__(self, t: float, u: TTrapezoidProfileState) -> None: ... def accept(self, t: float, u: TTrapezoidProfileState) -> None: ... + + +# Type Aliases +FloatSupplier: TypeAlias = Callable[[], float] +FloatOrFloatSupplier: TypeAlias = Union[float, Callable[[], float]]