[commands] Add onlyWhile and onlyIf (#5291)

wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/Command.java

@@ -82,9 +82,7 @@ public interface Command {
 
   /**
    * Decorates this command with an interrupt condition. If the specified condition becomes true
-   * before the command finishes normally, the command will be interrupted and un-scheduled. Note
-   * that this only applies to the command returned by this method; the calling command is not
-   * itself changed.
+   * before the command finishes normally, the command will be interrupted and un-scheduled.
    *
    * <p>Note: This decorator works by adding this command to a composition. The command the
    * decorator was called on cannot be scheduled independently or be added to a different
@@ -94,11 +92,30 @@ public interface Command {
    *
    * @param condition the interrupt condition
    * @return the command with the interrupt condition added
+   * @see #onlyWhile(BooleanSupplier)
    */
   default ParallelRaceGroup until(BooleanSupplier condition) {
     return raceWith(new WaitUntilCommand(condition));
   }
 
+  /**
+   * Decorates this command with a run condition. If the specified condition becomes false before
+   * the command finishes normally, the command will be interrupted and un-scheduled.
+   *
+   * <p>Note: This decorator works by adding this command to a composition. The command the
+   * decorator was called on cannot be scheduled independently or be added to a different
+   * composition (namely, decorators), unless it is manually cleared from the list of composed
+   * commands with {@link CommandScheduler#removeComposedCommand(Command)}. The command composition
+   * returned from this method can be further decorated without issue.
+   *
+   * @param condition the interrupt condition
+   * @return the command with the interrupt condition added
+   * @see #until(BooleanSupplier)
+   */
+  default ParallelRaceGroup onlyWhile(BooleanSupplier condition) {
+    return until(() -> !condition.getAsBoolean());
+  }
+
   /**
    * Decorates this command with an interrupt condition. If the specified condition becomes true
    * before the command finishes normally, the command will be interrupted and un-scheduled. Note
@@ -301,13 +318,39 @@ public interface Command {
    * running and the condition changes to true, the command will not stop running. The requirements
    * of this command will be kept for the new conditional command.
    *
+   * <p>Note: This decorator works by adding this command to a composition. The command the
+   * decorator was called on cannot be scheduled independently or be added to a different
+   * composition (namely, decorators), unless it is manually cleared from the list of composed
+   * commands with {@link CommandScheduler#removeComposedCommand(Command)}. The command composition
+   * returned from this method can be further decorated without issue.
+   *
    * @param condition the condition that will prevent the command from running
    * @return the decorated command
+   * @see #onlyIf(BooleanSupplier)
    */
   default ConditionalCommand unless(BooleanSupplier condition) {
     return new ConditionalCommand(new InstantCommand(), this, condition);
   }
 
+  /**
+   * Decorates this command to only run if this condition is met. If the command is already running
+   * and the condition changes to false, the command will not stop running. The requirements of this
+   * command will be kept for the new conditional command.
+   *
+   * <p>Note: This decorator works by adding this command to a composition. The command the
+   * decorator was called on cannot be scheduled independently or be added to a different
+   * composition (namely, decorators), unless it is manually cleared from the list of composed
+   * commands with {@link CommandScheduler#removeComposedCommand(Command)}. The command composition
+   * returned from this method can be further decorated without issue.
+   *
+   * @param condition the condition that will allow the command to run
+   * @return the decorated command
+   * @see #unless(BooleanSupplier)
+   */
+  default ConditionalCommand onlyIf(BooleanSupplier condition) {
+    return unless(() -> !condition.getAsBoolean());
+  }
+
   /**
    * Decorates this command to run or stop when disabled.
    *

wpilibNewCommands/src/main/native/cpp/frc2/command/Command.cpp

@@ -41,6 +41,10 @@ CommandPtr Command::Until(std::function<bool()> condition) && {
   return std::move(*this).ToPtr().Until(std::move(condition));
 }
 
+CommandPtr Command::OnlyWhile(std::function<bool()> condition) && {
+  return std::move(*this).ToPtr().OnlyWhile(std::move(condition));
+}
+
 CommandPtr Command::IgnoringDisable(bool doesRunWhenDisabled) && {
   return std::move(*this).ToPtr().IgnoringDisable(doesRunWhenDisabled);
 }
@@ -96,6 +100,10 @@ CommandPtr Command::Unless(std::function<bool()> condition) && {
   return std::move(*this).ToPtr().Unless(std::move(condition));
 }
 
+CommandPtr Command::OnlyIf(std::function<bool()> condition) && {
+  return std::move(*this).ToPtr().OnlyIf(std::move(condition));
+}
+
 CommandPtr Command::FinallyDo(std::function<void(bool)> end) && {
   return std::move(*this).ToPtr().FinallyDo(std::move(end));
 }

wpilibNewCommands/src/main/native/cpp/frc2/command/CommandPtr.cpp

@@ -151,6 +151,11 @@ CommandPtr CommandPtr::Until(std::function<bool()> condition) && {
   return std::move(*this);
 }
 
+CommandPtr CommandPtr::OnlyWhile(std::function<bool()> condition) && {
+  AssertValid();
+  return std::move(*this).Until(std::not_fn(std::move(condition)));
+}
+
 CommandPtr CommandPtr::Unless(std::function<bool()> condition) && {
   AssertValid();
   m_ptr = std::make_unique<ConditionalCommand>(
@@ -159,6 +164,11 @@ CommandPtr CommandPtr::Unless(std::function<bool()> condition) && {
   return std::move(*this);
 }
 
+CommandPtr CommandPtr::OnlyIf(std::function<bool()> condition) && {
+  AssertValid();
+  return std::move(*this).Unless(std::not_fn(std::move(condition)));
+}
+
 CommandPtr CommandPtr::DeadlineWith(CommandPtr&& parallel) && {
   AssertValid();
   std::vector<std::unique_ptr<Command>> vec;

wpilibNewCommands/src/main/native/include/frc2/command/Command.h

@@ -137,6 +137,17 @@ class Command {
    */
   [[nodiscard]] CommandPtr Until(std::function<bool()> condition) &&;
 
+  /**
+   * Decorates this command with a run condition.  If the specified condition
+   * becomes false before the command finishes normally, the command will be
+   * interrupted and un-scheduled. Note that this only applies to the command
+   * returned by this method; the calling command is not itself changed.
+   *
+   * @param condition the interrupt condition
+   * @return the command with the interrupt condition added
+   */
+  [[nodiscard]] CommandPtr OnlyWhile(std::function<bool()> condition) &&;
+
   /**
    * Decorates this command with an interrupt condition.  If the specified
    * condition becomes true before the command finishes normally, the command
@@ -245,6 +256,17 @@ safe) semantics.
    */
   [[nodiscard]] CommandPtr Unless(std::function<bool()> condition) &&;
 
+  /**
+   * Decorates this command to only run if this condition is met. If the command
+   * is already running and the condition changes to false, the command will not
+   * stop running. The requirements of this command will be kept for the new
+   * conditional command.
+   *
+   * @param condition the condition that will allow the command to run
+   * @return the decorated command
+   */
+  [[nodiscard]] CommandPtr OnlyIf(std::function<bool()> condition) &&;
+
   /**
    * Decorates this command to run or stop when disabled.
    *

wpilibNewCommands/src/main/native/include/frc2/command/CommandPtr.h

@@ -159,6 +159,17 @@ class CommandPtr final {
    */
   [[nodiscard]] CommandPtr Until(std::function<bool()> condition) &&;
 
+  /**
+   * Decorates this command with a run condition.  If the specified condition
+   * becomes false before the command finishes normally, the command will be
+   * interrupted and un-scheduled. Note that this only applies to the command
+   * returned by this method; the calling command is not itself changed.
+   *
+   * @param condition the interrupt condition
+   * @return the command with the interrupt condition added
+   */
+  [[nodiscard]] CommandPtr OnlyWhile(std::function<bool()> condition) &&;
+
   /**
    * Decorates this command to only run if this condition is not met. If the
    * command is already running and the condition changes to true, the command
@@ -170,6 +181,17 @@ class CommandPtr final {
    */
   [[nodiscard]] CommandPtr Unless(std::function<bool()> condition) &&;
 
+  /**
+   * Decorates this command to only run if this condition is met. If the command
+   * is already running and the condition changes to false, the command will not
+   * stop running. The requirements of this command will be kept for the new
+   * conditional command.
+   *
+   * @param condition the condition that will allow the command to run
+   * @return the decorated command
+   */
+  [[nodiscard]] CommandPtr OnlyIf(std::function<bool()> condition) &&;
+
   /**
    * Decorates this command with a set of commands to run parallel to it, ending
    * when the calling command ends and interrupting all the others. Often more

wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandDecoratorTest.java

@@ -57,6 +57,22 @@ class CommandDecoratorTest extends CommandTestBase {
     }
   }
 
+  @Test
+  void onlyWhileTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicBoolean condition = new AtomicBoolean(true);
+
+      Command command = new WaitCommand(10).onlyWhile(condition::get);
+
+      scheduler.schedule(command);
+      scheduler.run();
+      assertTrue(scheduler.isScheduled(command));
+      condition.set(false);
+      scheduler.run();
+      assertFalse(scheduler.isScheduled(command));
+    }
+  }
+
   @Test
   void ignoringDisableTest() {
     try (CommandScheduler scheduler = new CommandScheduler()) {
@@ -221,6 +237,26 @@ class CommandDecoratorTest extends CommandTestBase {
     }
   }
 
+  @Test
+  void onlyIfTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicBoolean onlyIfCondition = new AtomicBoolean(false);
+      AtomicBoolean hasRunCondition = new AtomicBoolean(false);
+
+      Command command =
+          new InstantCommand(() -> hasRunCondition.set(true)).onlyIf(onlyIfCondition::get);
+
+      scheduler.schedule(command);
+      scheduler.run();
+      assertFalse(hasRunCondition.get());
+
+      onlyIfCondition.set(true);
+      scheduler.schedule(command);
+      scheduler.run();
+      assertTrue(hasRunCondition.get());
+    }
+  }
+
   @Test
   void finallyDoTest() {
     try (CommandScheduler scheduler = new CommandScheduler()) {

wpilibNewCommands/src/test/native/cpp/frc2/command/CommandDecoratorTest.cpp

@@ -54,6 +54,24 @@ TEST_F(CommandDecoratorTest, Until) {
   EXPECT_FALSE(scheduler.IsScheduled(command));
 }
 
+TEST_F(CommandDecoratorTest, OnlyWhile) {
+  CommandScheduler scheduler = GetScheduler();
+
+  bool run = true;
+
+  auto command = RunCommand([] {}, {}).OnlyWhile([&run] { return run; });
+
+  scheduler.Schedule(command);
+
+  scheduler.Run();
+  EXPECT_TRUE(scheduler.IsScheduled(command));
+
+  run = false;
+
+  scheduler.Run();
+  EXPECT_FALSE(scheduler.IsScheduled(command));
+}
+
 TEST_F(CommandDecoratorTest, IgnoringDisable) {
   CommandScheduler scheduler = GetScheduler();
 
@@ -140,6 +158,27 @@ TEST_F(CommandDecoratorTest, Unless) {
   EXPECT_TRUE(hasRun);
 }
 
+TEST_F(CommandDecoratorTest, OnlyIf) {
+  CommandScheduler scheduler = GetScheduler();
+
+  bool hasRun = false;
+  bool onlyIfBool = false;
+
+  auto command =
+      InstantCommand([&hasRun] { hasRun = true; }, {}).OnlyIf([&onlyIfBool] {
+        return onlyIfBool;
+      });
+
+  scheduler.Schedule(command);
+  scheduler.Run();
+  EXPECT_FALSE(hasRun);
+
+  onlyIfBool = true;
+  scheduler.Schedule(command);
+  scheduler.Run();
+  EXPECT_TRUE(hasRun);
+}
+
 TEST_F(CommandDecoratorTest, FinallyDo) {
   CommandScheduler scheduler = GetScheduler();
   int first = 0;