[commands] Add FinallyDo and HandleInterrupt decorators (#4412)

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

@@ -4,6 +4,9 @@
 
 package edu.wpi.first.wpilibj2.command;
 
+import static edu.wpi.first.wpilibj.util.ErrorMessages.requireNonNullParam;
+
+import edu.wpi.first.util.function.BooleanConsumer;
 import java.util.Set;
 import java.util.function.BooleanSupplier;
 
@@ -348,6 +351,42 @@ public interface Command {
     };
   }
 
+  /**
+   * Decorates this command with a lambda to call on interrupt or end, following the command's
+   * inherent {@link #end(boolean)} method.
+   *
+   * @param end a lambda accepting a boolean parameter specifying whether the command was
+   *     interrupted.
+   * @return the decorated command
+   */
+  default WrapperCommand finallyDo(BooleanConsumer end) {
+    requireNonNullParam(end, "end", "Command.finallyDo()");
+    return new WrapperCommand(this) {
+      @Override
+      public void end(boolean interrupted) {
+        super.end(interrupted);
+        end.accept(interrupted);
+      }
+    };
+  }
+
+  /**
+   * Decorates this command with a lambda to call on interrupt, following the command's inherent
+   * {@link #end(boolean)} method.
+   *
+   * @param handler a lambda to run when the command is interrupted
+   * @return the decorated command
+   */
+  default WrapperCommand handleInterrupt(Runnable handler) {
+    requireNonNullParam(handler, "handler", "Command.handleInterrupt()");
+    return finallyDo(
+        interrupted -> {
+          if (interrupted) {
+            handler.run();
+          }
+        });
+  }
+
   /** Schedules this command. */
   default void schedule() {
     CommandScheduler.getInstance().schedule(this);

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

@@ -108,6 +108,16 @@ CommandPtr Command::Unless(std::function<bool()> condition) && {
       .Unless(std::move(condition));
 }
 
+CommandPtr Command::FinallyDo(std::function<void(bool)> end) && {
+  return CommandPtr(std::move(*this).TransferOwnership())
+      .FinallyDo(std::move(end));
+}
+
+CommandPtr Command::HandleInterrupt(std::function<void(void)> handler) && {
+  return CommandPtr(std::move(*this).TransferOwnership())
+      .HandleInterrupt(std::move(handler));
+}
+
 void Command::Schedule() {
   CommandScheduler::GetInstance().Schedule(this);
 }

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

@@ -167,6 +167,37 @@ CommandPtr CommandPtr::RaceWith(CommandPtr&& parallel) && {
   return std::move(*this);
 }
 
+namespace {
+class FinallyCommand : public WrapperCommand {
+ public:
+  FinallyCommand(std::unique_ptr<Command>&& command,
+                 std::function<void(bool)> end)
+      : WrapperCommand(std::move(command)), m_end(std::move(end)) {}
+
+  void End(bool interrupted) override {
+    WrapperCommand::End(interrupted);
+    m_end(interrupted);
+  }
+
+ private:
+  std::function<void(bool)> m_end;
+};
+}  // namespace
+
+CommandPtr CommandPtr::FinallyDo(std::function<void(bool)> end) && {
+  m_ptr = std::make_unique<FinallyCommand>(std::move(m_ptr), std::move(end));
+  return std::move(*this);
+}
+
+CommandPtr CommandPtr::HandleInterrupt(std::function<void(void)> handler) && {
+  return std::move(*this).FinallyDo(
+      [handler = std::move(handler)](bool interrupted) {
+        if (interrupted) {
+          handler();
+        }
+      });
+}
+
 Command* CommandPtr::get() const {
   return m_ptr.get();
 }

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

@@ -261,6 +261,25 @@ class Command {
   [[nodiscard]] CommandPtr WithInterruptBehavior(
       Command::InterruptionBehavior interruptBehavior) &&;
 
+  /**
+   * Decorates this command with a lambda to call on interrupt or end, following
+   * the command's inherent Command::End(bool) method.
+   *
+   * @param end a lambda accepting a boolean parameter specifying whether the
+   * command was interrupted.
+   * @return the decorated command
+   */
+  [[nodiscard]] CommandPtr FinallyDo(std::function<void(bool)> end) &&;
+
+  /**
+   * Decorates this command with a lambda to call on interrupt, following the
+   * command's inherent Command::End(bool) method.
+   *
+   * @param handler a lambda to run when the command is interrupted
+   * @return the decorated command
+   */
+  [[nodiscard]] CommandPtr HandleInterrupt(std::function<void()> handler) &&;
+
   /**
    * Schedules this command.
    */

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

@@ -207,6 +207,25 @@ class CommandPtr final {
    */
   [[nodiscard]] CommandPtr RaceWith(CommandPtr&& parallel) &&;
 
+  /**
+   * Decorates this command with a lambda to call on interrupt or end, following
+   * the command's inherent Command::End(bool) method.
+   *
+   * @param end a lambda accepting a boolean parameter specifying whether the
+   * command was interrupted.
+   * @return the decorated command
+   */
+  [[nodiscard]] CommandPtr FinallyDo(std::function<void(bool)> end) &&;
+
+  /**
+   * Decorates this command with a lambda to call on interrupt, following the
+   * command's inherent Command::End(bool) method.
+   *
+   * @param handler a lambda to run when the command is interrupted
+   * @return the decorated command
+   */
+  [[nodiscard]] CommandPtr HandleInterrupt(std::function<void()> handler) &&;
+
   /**
    * Get a raw pointer to the held command.
    */