feat(core): Add helpers for mapping ArgumentParseResults (#745)

Add helpers for mapping ArgumentParseResults fixes #738 closes #739
Incendo · May 30, 2024 · 6f08ae2 · 6f08ae2
1 parent 611a55d
commit 6f08ae2
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 7 deletions.
diff --git a/cloud-core/src/main/java/org/incendo/cloud/parser/ArgumentParseResult.java b/cloud-core/src/main/java/org/incendo/cloud/parser/ArgumentParseResult.java
@@ -26,6 +26,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.function.Function;
 import org.apiguardian.api.API;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.incendo.cloud.exception.handling.ExceptionController;
@@ -115,6 +116,49 @@ private ArgumentParseResult() {
         return CompletableFuture.completedFuture(this);
     }
 
+    /**
+     * Returns a future resulting from applying {@code mapper} to the parsed value,
+     * or a completed future with the same failure as this result.
+     *
+     * @param mapper mapper
+     * @param <O>    new result value type
+     * @return new result future
+     */
+    public abstract <O> @NonNull CompletableFuture<ArgumentParseResult<O>> flatMapSuccessFuture(
+            @NonNull Function<T, CompletableFuture<ArgumentParseResult<O>>> mapper
+    );
+
+    /**
+     * Returns a success future resulting from applying {@code mapper} to the parsed value and
+     * wrapping in {@link #success(Object)}, or a completed future with the same failure as this result.
+     *
+     * @param mapper mapper
+     * @param <O>    new result value type
+     * @return new result future
+     */
+    public abstract <O> @NonNull CompletableFuture<ArgumentParseResult<O>> mapSuccessFuture(
+            @NonNull Function<T, CompletableFuture<O>> mapper
+    );
+
+    /**
+     * Returns the result from applying {@code mapper} to the parsed value,
+     * or the same failure as this result.
+     *
+     * @param mapper mapper
+     * @param <O>    new result value type
+     * @return new result
+     */
+    public abstract <O> @NonNull ArgumentParseResult<O> flatMapSuccess(@NonNull Function<T, ArgumentParseResult<O>> mapper);
+
+    /**
+     * Returns the result from applying {@code mapper} to the parsed value and
+     * wrapping in {@link #success(Object)}, or the same failure as this result.
+     *
+     * @param mapper mapper
+     * @param <O>    new result value type
+     * @return new result
+     */
+    public abstract <O> @NonNull ArgumentParseResult<O> mapSuccess(@NonNull Function<T, O> mapper);
 
     private static final class ParseSuccess<T> extends ArgumentParseResult<T> {
 
@@ -136,6 +180,32 @@ private ParseSuccess(final @NonNull T value) {
         public @NonNull Optional<Throwable> failure() {
             return Optional.empty();
         }
+
+        @Override
+        public <O> @NonNull CompletableFuture<ArgumentParseResult<O>> flatMapSuccessFuture(
+                final @NonNull Function<T, CompletableFuture<ArgumentParseResult<O>>> mapper
+        ) {
+            return mapper.apply(this.value);
+        }
+
+        @Override
+        public <O> @NonNull CompletableFuture<ArgumentParseResult<O>> mapSuccessFuture(
+                final @NonNull Function<T, CompletableFuture<O>> mapper
+        ) {
+            return mapper.apply(this.value).thenApply(ArgumentParseResult::success);
+        }
+
+        @Override
+        public <O> @NonNull ArgumentParseResult<O> flatMapSuccess(
+                final @NonNull Function<T, ArgumentParseResult<O>> mapper
+        ) {
+            return mapper.apply(this.value);
+        }
+
+        @Override
+        public <O> @NonNull ArgumentParseResult<O> mapSuccess(final @NonNull Function<T, O> mapper) {
+            return ArgumentParseResult.success(mapper.apply(this.value));
+        }
     }
 
     private static final class ParseFailure<T> extends ArgumentParseResult<T> {
@@ -158,5 +228,34 @@ private ParseFailure(final @NonNull Throwable failure) {
         public @NonNull Optional<Throwable> failure() {
             return Optional.of(this.failure);
         }
+
+        @Override
+        public <O> @NonNull CompletableFuture<ArgumentParseResult<O>> flatMapSuccessFuture(
+                final @NonNull Function<T, CompletableFuture<ArgumentParseResult<O>>> mapper
+        ) {
+            return CompletableFuture.completedFuture(this.self());
+        }
+
+        @Override
+        public <O> @NonNull CompletableFuture<ArgumentParseResult<O>> mapSuccessFuture(
+                final @NonNull Function<T, CompletableFuture<O>> mapper
+        ) {
+            return CompletableFuture.completedFuture(this.self());
+        }
+
+        @Override
+        public <O> @NonNull ArgumentParseResult<O> flatMapSuccess(final @NonNull Function<T, ArgumentParseResult<O>> mapper) {
+            return this.self();
+        }
+
+        @Override
+        public <O> @NonNull ArgumentParseResult<O> mapSuccess(final @NonNull Function<T, O> mapper) {
+            return this.self();
+        }
+
+        @SuppressWarnings("unchecked")
+        private <O> @NonNull ArgumentParseResult<O> self() {
+            return (ArgumentParseResult<O>) this;
+        }
     }
 }
diff --git a/cloud-core/src/main/java/org/incendo/cloud/parser/ArgumentParser.java b/cloud-core/src/main/java/org/incendo/cloud/parser/ArgumentParser.java
@@ -133,12 +133,7 @@ public interface ArgumentParser<C, T> extends SuggestionProviderHolder<C> {
             final @NonNull BiFunction<CommandContext<C>, T, CompletableFuture<ArgumentParseResult<O>>> mapper
     ) {
         requireNonNull(mapper, "mapper");
-        return this.flatMap((ctx, orig) -> {
-            if (orig.failure().isPresent()) {
-                return ArgumentParseResult.failureFuture(orig.failure().get());
-            }
-            return mapper.apply(ctx, orig.parsedValue().get());
-        });
+        return this.flatMap((ctx, result) -> result.flatMapSuccessFuture(value -> mapper.apply(ctx, value)));
     }
 
     /**
@@ -154,7 +149,7 @@ public interface ArgumentParser<C, T> extends SuggestionProviderHolder<C> {
             final @NonNull BiFunction<CommandContext<C>, T, CompletableFuture<O>> mapper
     ) {
         requireNonNull(mapper, "mapper");
-        return this.flatMapSuccess((ctx, orig) -> mapper.apply(ctx, orig).thenApply(ArgumentParseResult::success));
+        return this.flatMap((ctx, result) -> result.mapSuccessFuture(value -> mapper.apply(ctx, value)));
     }
 
     /**