From 6fea807991eec3376c2ccdb1d1f8a50a2c2ebd5c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20Spie=C3=9F?= <business@minn.dev>
Date: Mon, 26 Feb 2024 18:51:30 +0100
Subject: [PATCH] Allow map and delay operators to set reason

---
 .../restaction/AuditableRestAction.java       | 96 ++++++++++++++++++-
 .../operator/RestActionOperator.java          | 16 +++-
 2 files changed, 108 insertions(+), 4 deletions(-)

diff --git a/src/main/java/net/dv8tion/jda/api/requests/restaction/AuditableRestAction.java b/src/main/java/net/dv8tion/jda/api/requests/restaction/AuditableRestAction.java
index f595ffa49c..693deb391f 100644
--- a/src/main/java/net/dv8tion/jda/api/requests/restaction/AuditableRestAction.java
+++ b/src/main/java/net/dv8tion/jda/api/requests/restaction/AuditableRestAction.java
@@ -21,11 +21,14 @@
 import net.dv8tion.jda.api.entities.UserSnowflake;
 import net.dv8tion.jda.api.requests.RestAction;
 import net.dv8tion.jda.api.requests.restaction.pagination.AuditLogPaginationAction;
+import net.dv8tion.jda.api.utils.Result;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import java.time.Duration;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
-import java.util.function.BooleanSupplier;
+import java.util.function.*;
 
 /**
  * Extension of RestAction to allow setting a reason.
@@ -94,4 +97,95 @@ default AuditableRestAction<T> deadline(long timestamp)
     {
         return (AuditableRestAction<T>) RestAction.super.deadline(timestamp);
     }
+
+    @Nonnull
+    @Override
+    default <O> AuditableRestAction<O> map(@Nonnull Function<? super T, ? extends O> map)
+    {
+        return (AuditableRestAction<O>) RestAction.super.map(map);
+    }
+
+    @Nonnull
+    @Override
+    default <O> AuditableRestAction<O> flatMap(@Nonnull Function<? super T, ? extends RestAction<O>> flatMap)
+    {
+        return (AuditableRestAction<O>) RestAction.super.flatMap(flatMap);
+    }
+
+    @Nonnull
+    @Override
+    default <O> AuditableRestAction<O> flatMap(@Nullable Predicate<? super T> condition, @Nonnull Function<? super T, ? extends RestAction<O>> flatMap)
+    {
+        return (AuditableRestAction<O>) RestAction.super.flatMap(condition, flatMap);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<Result<T>> mapToResult()
+    {
+        return (AuditableRestAction<Result<T>>) RestAction.super.mapToResult();
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> onSuccess(@Nonnull Consumer<? super T> consumer)
+    {
+        return (AuditableRestAction<T>) RestAction.super.onSuccess(consumer);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> onErrorMap(@Nonnull Function<? super Throwable, ? extends T> map)
+    {
+        return (AuditableRestAction<T>) RestAction.super.onErrorMap(map);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> onErrorMap(@Nullable Predicate<? super Throwable> condition, @Nonnull Function<? super Throwable, ? extends T> map)
+    {
+        return (AuditableRestAction<T>) RestAction.super.onErrorMap(condition, map);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> onErrorFlatMap(@Nonnull Function<? super Throwable, ? extends RestAction<? extends T>> map)
+    {
+        return (AuditableRestAction<T>) RestAction.super.onErrorFlatMap(map);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> onErrorFlatMap(@Nullable Predicate<? super Throwable> condition, @Nonnull Function<? super Throwable, ? extends RestAction<? extends T>> map)
+    {
+        return (AuditableRestAction<T>) RestAction.super.onErrorFlatMap(condition, map);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> delay(@Nonnull Duration duration)
+    {
+        return (AuditableRestAction<T>) RestAction.super.delay(duration);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> delay(@Nonnull Duration duration, @Nullable ScheduledExecutorService scheduler)
+    {
+        return (AuditableRestAction<T>) RestAction.super.delay(duration, scheduler);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> delay(long delay, @Nonnull TimeUnit unit)
+    {
+        return (AuditableRestAction<T>) RestAction.super.delay(delay, unit);
+    }
+
+    @Nonnull
+    @Override
+    default AuditableRestAction<T> delay(long delay, @Nonnull TimeUnit unit, @Nullable ScheduledExecutorService scheduler)
+    {
+        return (AuditableRestAction<T>) RestAction.super.delay(delay, unit, scheduler);
+    }
 }
diff --git a/src/main/java/net/dv8tion/jda/internal/requests/restaction/operator/RestActionOperator.java b/src/main/java/net/dv8tion/jda/internal/requests/restaction/operator/RestActionOperator.java
index 83099c5c18..25f7228df2 100644
--- a/src/main/java/net/dv8tion/jda/internal/requests/restaction/operator/RestActionOperator.java
+++ b/src/main/java/net/dv8tion/jda/internal/requests/restaction/operator/RestActionOperator.java
@@ -19,13 +19,14 @@
 import net.dv8tion.jda.api.JDA;
 import net.dv8tion.jda.api.exceptions.ContextException;
 import net.dv8tion.jda.api.requests.RestAction;
+import net.dv8tion.jda.api.requests.restaction.AuditableRestAction;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 
-public abstract class RestActionOperator<I, O> implements RestAction<O>
+public abstract class RestActionOperator<I, O> implements AuditableRestAction<O>
 {
     protected BooleanSupplier check;
     protected long deadline = -1;
@@ -70,6 +71,15 @@ protected void handle(RestAction<I> action, Consumer<? super Throwable> failure,
         }, catcher);
     }
 
+    @Nonnull
+    @Override
+    public AuditableRestAction<O> reason(@Nullable String reason)
+    {
+        if (action instanceof AuditableRestAction)
+            ((AuditableRestAction<I>) action).reason(reason);
+        return this;
+    }
+
     @Nonnull
     @Override
     public JDA getJDA()
@@ -79,7 +89,7 @@ public JDA getJDA()
 
     @Nonnull
     @Override
-    public RestAction<O> setCheck(@Nullable BooleanSupplier checks)
+    public AuditableRestAction<O> setCheck(@Nullable BooleanSupplier checks)
     {
         this.check = checks;
         action.setCheck(checks);
@@ -95,7 +105,7 @@ public BooleanSupplier getCheck()
 
     @Nonnull
     @Override
-    public RestAction<O> deadline(long timestamp)
+    public AuditableRestAction<O> deadline(long timestamp)
     {
         this.deadline = timestamp;
         action.deadline(timestamp);