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,6 @@
+from ..trigger import Trigger
+
+
+class Button(Trigger):
+    whenPressed = Trigger.whenActive
+    whenReleased = Trigger.whenInactive

commands2/button/joystickbutton.py

@@ -0,0 +1,8 @@
+from wpilib import Joystick
+
+from .button import Button
+
+
+class JoystickButton(Button):
+    def __init__(self, joystick: Joystick, button: int) -> None:
+        super().__init__(lambda: joystick.getRawButton(button))

commands2/button/networkbutton.py

@@ -0,0 +1,32 @@
+from networktables import NetworkTable, NetworkTables, NetworkTableEntry
+
+from typing import Union, overload
+
+from .button import Button
+
+
+class NetworkButton(Button):
+    @overload
+    def __init__(self, entry: NetworkTableEntry) -> None:
+        ...
+
+    @overload
+    def __init__(self, table: Union[NetworkTable, str], field: str) -> None:
+        ...
+
+    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)
+            )
+        else:
+            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)

commands2/button/povbutton.py

@@ -0,0 +1,8 @@
+from wpilib import Joystick
+
+from .button import Button
+
+
+class POVButton(Button):
+    def __init__(self, joystick: Joystick, angle: int, povNumber: int = 0) -> None:
+        super().__init__(lambda: joystick.getPOV(povNumber) == angle)

commands2/coroutinecommand.py

@@ -0,0 +1,96 @@
+from functools import wraps
+from typing import Any, Callable, Generator, List, Union, Optional
+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):
+    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:
+        self.coroutine = None
+        self.coroutine_function = None
+        self.runsWhenDisabled = lambda: runs_when_disabled
+
+        if is_coroutine(coroutine):
+            self.coroutine = coroutine
+        elif is_coroutine_function(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:
+            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
+
+
+class commandify:
+    def __init__(self, requirements: Optional[List[Subsystem]] = None) -> None:
+        self.requirements = requirements
+
+    def __call__(self, func: Coroutineable):
+        @wraps(func)
+        def arg_accepter(*args, **kwargs) -> CoroutineCommand:
+            return CoroutineCommand(
+                lambda: ensure_generator_function(func)(*args, **kwargs),
+                self.requirements,
+            )
+
+        return arg_accepter

commands2/src/Command.cpp.inl

@@ -159,5 +159,12 @@ 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>(),
+    "qweqwe\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&);

commands2/trigger.py

@@ -0,0 +1,149 @@
+from typing import Callable, Optional, overload, List, Union
+
+from ._impl import Command, Subsystem
+from ._impl import Trigger as _Trigger
+
+from .coroutinecommand import CoroutineCommand, Coroutineable, Coroutine
+
+
+class Trigger:
+    """
+    A button that can be pressed or released.
+    """
+
+    def __init__(self, is_active: Callable[[], bool] = lambda: False) -> None:
+        self._trigger = _Trigger(is_active)
+
+    def __bool__(self) -> bool:
+        return bool(self._trigger)
+
+    def get(self) -> bool:
+        return bool(self)
+
+    def __call__(self) -> bool:
+        return bool(self)
+
+    def __and__(self, other: "Trigger") -> "Trigger":
+        return Trigger(lambda: self() and other())
+
+    def __or__(self, other: "Trigger") -> "Trigger":
+        return Trigger(lambda: self() or other())
+
+    def __not__(self) -> "Trigger":
+        return Trigger(lambda: not self())
+
+    @overload
+    def whenActive(self, command: Command, /, interruptible: bool = True) -> None:
+        ...
+
+    @overload
+    def whenActive(
+        self,
+        coroutine: Union[Coroutine, Coroutineable],
+        /,
+        *,
+        interruptible: bool = True,
+        requirements: Optional[List[Subsystem]] = None,
+        runs_when_disabled: bool = False,
+    ) -> None:
+        ...
+
+    @overload
+    def whenActive(
+        self,
+        coroutine: None,
+        /,
+        *,
+        interruptible: bool = True,
+        requirements: Optional[List[Subsystem]] = None,
+        runs_when_disabled: bool = False,
+    ) -> Callable[[Coroutineable], None]:
+        ...
+
+    def whenActive(
+        self,
+        command_or_coroutine: Optional[Union[Command, Coroutine, Coroutineable]],
+        /,
+        interruptible: bool = True,
+        requirements: Optional[List[Subsystem]] = None,
+        runs_when_disabled: bool = False,
+    ) -> Union[None, Callable[[Coroutineable], None]]:
+        if command_or_coroutine is None:
+
+            def wrapper(coroutine: Coroutineable) -> None:
+                self.whenActive(
+                    coroutine,
+                    interruptible=interruptible,
+                    requirements=requirements,
+                    runs_when_disabled=runs_when_disabled,
+                )
+
+            return wrapper
+
+        if isinstance(command_or_coroutine, Command):
+            self._trigger.whenActive(command_or_coroutine, interruptible)
+            return
+
+        self._trigger.whenActive(
+            CoroutineCommand(command_or_coroutine, requirements, runs_when_disabled),
+            interruptible,
+        )
+        return
+
+    @overload
+    def whenInactive(self, command: Command, /, interruptible: bool = True) -> None:
+        ...
+
+    @overload
+    def whenInactive(
+        self,
+        coroutine: Union[Coroutine, Coroutineable],
+        /,
+        *,
+        interruptible: bool = True,
+        requirements: Optional[List[Subsystem]] = None,
+        runs_when_disabled: bool = False,
+    ) -> None:
+        ...
+
+    @overload
+    def whenInactive(
+        self,
+        coroutine: None,
+        /,
+        *,
+        interruptible: bool = True,
+        requirements: Optional[List[Subsystem]] = None,
+        runs_when_disabled: bool = False,
+    ) -> Callable[[Coroutineable], None]:
+        ...
+
+    def whenInactive(
+        self,
+        command_or_coroutine: Optional[Union[Command, Coroutine, Coroutineable]],
+        /,
+        interruptible: bool = True,
+        requirements: Optional[List[Subsystem]] = None,
+        runs_when_disabled: bool = False,
+    ) -> Union[None, Callable[[Coroutineable], None]]:
+        if command_or_coroutine is None:
+
+            def wrapper(coroutine: Coroutineable) -> None:
+                self.whenInactive(
+                    coroutine,
+                    interruptible=interruptible,
+                    requirements=requirements,
+                    runs_when_disabled=runs_when_disabled,
+                )
+
+            return wrapper
+
+        if isinstance(command_or_coroutine, Command):
+            self._trigger.whenInactive(command_or_coroutine, interruptible)
+            return
+
+        self._trigger.whenInactive(
+            CoroutineCommand(command_or_coroutine, requirements, runs_when_disabled),
+            interruptible,
+        )
+        return