Add coroutine command support

commands2/__init__.py

@@ -1,4 +1,6 @@
 from . import _init_impl
+from .trigger import Trigger
+from .coroutinecommand import CoroutineCommand, commandify
 
 from .version import version as __version__
 
@@ -40,7 +42,7 @@
     TrapezoidProfileCommandRadians,
     TrapezoidProfileSubsystem,
     TrapezoidProfileSubsystemRadians,
-    Trigger,
+    # Trigger,
     WaitCommand,
     WaitUntilCommand,
     # button,
@@ -84,9 +86,12 @@
     "TrapezoidProfileCommandRadians",
     "TrapezoidProfileSubsystem",
     "TrapezoidProfileSubsystemRadians",
-    "Trigger",
+    # "_Trigger",
     "WaitCommand",
     "WaitUntilCommand",
     # "button",
     "requirementsDisjoint",
+    "commandify",
+    "CoroutineCommand",
+    "Trigger",
 ]

commands2/button/__init__.py

@@ -1,4 +1,6 @@
-# autogenerated by 'robotpy-build create-imports commands2.button commands2._impl.button'
-from .._impl.button import Button, JoystickButton, NetworkButton, POVButton
+from .button import Button
+from .joystickbutton import JoystickButton
+from .networkbutton import NetworkButton
+from .povbutton import POVButton
 
 __all__ = ["Button", "JoystickButton", "NetworkButton", "POVButton"]

commands2/button/button.py

@@ -0,0 +1,17 @@
+from ..trigger import Trigger
+
+
+class Button(Trigger):
+    """
+    A class used to bind command scheduling to button presses.
+    Can be composed with other buttons with the operators in Trigger.
+
+    @see Trigger
+    """
+
+    whenPressed = Trigger.whenActive
+    whenReleased = Trigger.whenInactive
+    whileHeld = Trigger.whileActiveContinous
+    whenHeld = Trigger.whileActiveOnce
+    toggleWhenPressed = Trigger.toggleWhenActive
+    cancelWhenPressed = Trigger.cancelWhenActive

commands2/button/joystickbutton.py

@@ -0,0 +1,21 @@
+from wpilib import Joystick
+
+from .button import Button
+
+
+class JoystickButton(Button):
+    """
+    A class used to bind command scheduling to joystick button presses.
+    Can be composed with other buttons with the operators in Trigger.
+
+    @see Trigger
+    """
+
+    def __init__(self, joystick: Joystick, button: int) -> None:
+        """
+        Creates a JoystickButton that commands can be bound to.
+
+        :param joystick: The joystick on which the button is located.
+        :param button: The number of the button on the joystick.
+        """
+        super().__init__(lambda: joystick.getRawButton(button))

commands2/button/networkbutton.py

@@ -0,0 +1,52 @@
+from networktables import NetworkTable, NetworkTables, NetworkTableEntry
+
+from typing import Union, overload
+
+from .button import Button
+
+
+class NetworkButton(Button):
+    """
+    A class used to bind command scheduling to a NetworkTable boolean fields.
+    Can be composed with other buttons with the operators in Trigger.
+
+    @see Trigger
+    """
+
+    @overload
+    def __init__(self, entry: NetworkTableEntry) -> None:
+        """
+        Creates a NetworkButton that commands can be bound to.
+
+        :param entry: The entry that is the value.
+        """
+
+    @overload
+    def __init__(self, table: Union[NetworkTable, str], field: str) -> None:
+        """
+        Creates a NetworkButton that commands can be bound to.
+
+        :param table: The table where the networktable value is located.
+        :param field: The field that is the value.
+        """
+
+    def __init__(self, *args, **kwargs) -> None:
+        num_args = len(args) + len(kwargs)
+        if num_args == 1:
+            entry: NetworkTableEntry = kwargs.get("entry", args[0])
+            super().__init__(
+                lambda: NetworkTables.isConnected() and entry.getBoolean(False)
+            )
+        elif num_args == 2:
+            table = kwargs.get("table", args[0])
+            field = kwargs.get("field", args[-1])
+
+            if isinstance(table, str):
+                table = NetworkTables.getTable(table)
+
+            entry = table.getEntry(field)
+            self.__init__(entry)
+        else:
+            raise TypeError(
+                f"__init__() takes 1 or 2 positional arguments but {num_args} were given"
+            )

commands2/button/povbutton.py

@@ -0,0 +1,22 @@
+from wpilib import Joystick
+
+from .button import Button
+
+
+class POVButton(Button):
+    """
+    A class used to bind command scheduling to joystick POV presses.
+    Can be composed with other buttons with the operators in Trigger.
+
+    @see Trigger
+    """
+
+    def __init__(self, joystick: Joystick, angle: int, povNumber: int = 0) -> None:
+        """
+        Creates a POVButton that commands can be bound to.
+
+        :param joystick: The joystick on which the button is located.
+        :param angle: The angle of the POV corresponding to a button press.
+        :param povNumber: The number of the POV on the joystick.
+        """
+        super().__init__(lambda: joystick.getPOV(povNumber) == angle)

commands2/coroutinecommand.py

@@ -0,0 +1,139 @@
+from functools import wraps
+from typing import Any, Callable, Generator, List, Union, Optional, overload
+from ._impl import CommandBase, Subsystem
+import inspect
+from typing_extensions import TypeGuard
+
+Coroutine = Generator[None, None, None]
+CoroutineFunction = Callable[[], Generator[None, None, None]]
+Coroutineable = Union[Callable[[], None], CoroutineFunction]
+
+
+def is_coroutine(func: Any) -> TypeGuard[Coroutine]:
+    return inspect.isgenerator(func)
+
+
+def is_coroutine_function(func: Any) -> TypeGuard[CoroutineFunction]:
+    return inspect.isgeneratorfunction(func)
+
+
+def is_coroutineable(func: Any) -> TypeGuard[Coroutineable]:
+    return is_coroutine_function(func) or callable(func)
+
+
+def ensure_generator_function(func: Coroutineable) -> Callable[..., Coroutine]:
+    if is_coroutine_function(func):
+        return func
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        func(*args, **kwargs)
+        yield
+
+    return wrapper
+
+
+class CoroutineCommand(CommandBase):
+    """
+    A class that wraps a coroutine function into a command.
+    """
+
+    coroutine: Optional[Coroutine]
+    coroutine_function: Optional[Coroutineable]
+    is_finished: bool
+
+    def __init__(
+        self,
+        coroutine: Union[Coroutine, Coroutineable],
+        requirements: Optional[List[Subsystem]] = None,
+        runs_when_disabled: bool = False,
+    ) -> None:
+        """
+        Creates a CoroutineCommand than can be used as a command.
+
+        :param coroutine: The coroutine or coroutine function to bind.
+        :param requirements: The subsystems that this command requires.
+        :param runs_when_disabled: Whether or not this command runs when the robot is disabled.
+        """
+        self.coroutine = None
+        self.coroutine_function = None
+        self.runsWhenDisabled = lambda: runs_when_disabled
+
+        if is_coroutine(coroutine):
+            self.coroutine = coroutine
+        elif is_coroutineable(coroutine):
+            self.coroutine_function = coroutine
+        else:
+            raise TypeError("The coroutine must be a coroutine or a coroutine function")
+
+        if requirements is not None:
+            self.addRequirements(requirements)
+
+        self.is_finished = False
+
+    def initialize(self) -> None:
+        if self.coroutine_function:
+            self.coroutine = ensure_generator_function(self.coroutine_function)()
+        elif self.coroutine and self.is_finished:
+            raise RuntimeError("Generator objects cannot be reused.")
+
+        self.is_finished = False
+
+    def execute(self):
+        try:
+            if not self.is_finished:
+                if not self.coroutine:
+                    raise TypeError("This command was not properly initialized")
+                next(self.coroutine)
+        except StopIteration:
+            self.is_finished = True
+
+    def isFinished(self):
+        return self.is_finished
+
+
+@overload
+def commandify(
+    *, requirements: Optional[List[Subsystem]] = None, runs_when_disabled: bool = False
+) -> Callable[[Coroutineable], Callable[..., CoroutineCommand]]:
+    """
+    A decorator that turns a coroutine function into a command.
+    A def should be under this.
+
+    :param requirements: The subsystems that this command requires.
+    :param runs_when_disabled: Whether or not this command runs when the robot is disabled.
+    """
+
+
+@overload
+def commandify(coroutine: Coroutineable, /) -> Callable[..., CoroutineCommand]:
+    """
+    A decorator that turns a coroutine function into a command.
+    A def should be under this.
+    """
+
+
+def commandify(
+    coroutine: Optional[Coroutineable] = None,
+    /,
+    *,
+    requirements: Optional[List[Subsystem]] = None,
+    runs_when_disabled: bool = False,
+) -> Union[
+    Callable[[Coroutineable], Callable[..., CoroutineCommand]],
+    Callable[..., CoroutineCommand],
+]:
+    def wrapper(func: Coroutineable) -> Callable[..., CoroutineCommand]:
+        @wraps(func)
+        def arg_accepter(*args, **kwargs) -> CoroutineCommand:
+            return CoroutineCommand(
+                lambda: ensure_generator_function(func)(*args, **kwargs),
+                requirements,
+            )
+
+        return arg_accepter
+
+    if coroutine is None:
+        return wrapper
+
+    return wrapper(coroutine)

commands2/src/Command.cpp.inl

@@ -159,5 +159,17 @@ cls_Command
     "\n"
     ":returns: the command with the timeout added\n"
     DECORATOR_NOTE)
+  .def("__iter__",
+    [](std::shared_ptr<Command> self) {
+      return py::make_iterator(CommandIterator(self), CommandIteratorSentinel());
+    },
+    py::keep_alive<0, 1>(),
+    "Creates an Iterator for this command. The iterator will run the command and\n"
+    "will only exhaust when the command is finished.\n"
+    "Note that the iterator will not run the command in the background. It must be\n"
+    "explicitly be iterated over.\n"
+    "\n"
+    ":returns: an iterator for this command\n"
+    )
   ;
 

commands2/src/helpers.cpp

@@ -23,4 +23,20 @@ std::vector<frc2::Subsystem*> pyargs2SubsystemList(py::args subs) {
     subsystems.emplace_back(py::cast<frc2::Subsystem*>(sub));
   }
   return subsystems;
-}
+}
+
+CommandIterator::CommandIterator(std::shared_ptr<frc2::Command> cmd) : cmd(cmd) {}
+std::shared_ptr<frc2::Command> CommandIterator::operator*() const { return cmd; }
+CommandIterator& CommandIterator::operator++() {
+  if (!called_initialize) {
+    cmd->Initialize();
+    called_initialize = true;
+    return *this;
+  }
+  cmd->Execute();
+  return *this;
+}
+
+bool operator==(const CommandIterator& it, const CommandIteratorSentinel&) {
+  return it.called_initialize && it.cmd->IsFinished();
+}

commands2/src/helpers.h

@@ -16,4 +16,17 @@ std::shared_ptr<T> convertToSharedPtrHack(T *orig) {
 
     py::object pyo = py::cast(orig);
     return py::cast<std::shared_ptr<T>>(pyo);
-}
+}
+
+class CommandIterator {
+    public:
+        std::shared_ptr<frc2::Command> cmd;
+        bool called_initialize = false;
+        explicit CommandIterator(std::shared_ptr<frc2::Command> cmd);
+        std::shared_ptr<frc2::Command> operator*() const;
+        CommandIterator& operator++();
+};
+
+class CommandIteratorSentinel {};
+
+bool operator==(const CommandIterator& it, const CommandIteratorSentinel&);