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

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.
    */