Improve mapping/filtering

src/main/java/net/minestom/server/command/builder/arguments/Argument.java

@@ -11,6 +11,8 @@ import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 
 /**
@@ -260,10 +262,21 @@ public abstract class Argument<T> {
      * @return A new ArgumentMap that can get this complex object type.
      */
     @ApiStatus.Experimental
-    public <O> @NotNull ArgumentMap<T, O> map(@NotNull ArgumentMap.Mapper<T, O> mapper) {
+    public <O> @NotNull Argument<O> map(@NotNull Function<T, O> mapper) {
         return new ArgumentMap<>(this, mapper);
     }
 
+    /**
+     * Maps this argument's output to another result.
+     *
+     * @param predicate the argument predicate
+     * @return A new ArgumentMap that filters using this filterer.
+     */
+    @ApiStatus.Experimental
+    public @NotNull Argument<T> filter(@NotNull Predicate<T> predicate) {
+        return new ArgumentFilter<>(this, predicate);
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -278,4 +291,63 @@ public abstract class Argument<T> {
     public int hashCode() {
         return id.hashCode();
     }
+
+    private static final class ArgumentMap<I, O> extends Argument<O> {
+        public static final int INVALID_MAP = 555;
+        final Argument<I> argument;
+        final Function<I, O> mapper;
+
+        private ArgumentMap(@NotNull Argument<I> argument, @NotNull Function<I, O> mapper) {
+            super(argument.getId(), argument.allowSpace(), argument.useRemaining());
+            if (argument.getSuggestionCallback() != null)
+                this.setSuggestionCallback(argument.getSuggestionCallback());
+            if (argument.getDefaultValue() != null)
+                this.setDefaultValue(() -> mapper.apply(argument.getDefaultValue().get()));
+            this.argument = argument;
+            this.mapper = mapper;
+        }
+
+        @Override
+        public @NotNull O parse(@NotNull String input) throws ArgumentSyntaxException {
+            final I value = argument.parse(input);
+            final O mappedValue = mapper.apply(value);
+            if (mappedValue == null)
+                throw new ArgumentSyntaxException("Couldn't be converted to map type", input, INVALID_MAP);
+            return mappedValue;
+        }
+
+        @Override
+        public void processNodes(@NotNull NodeMaker nodeMaker, boolean executable) {
+            argument.processNodes(nodeMaker, executable);
+        }
+    }
+
+    private static final class ArgumentFilter<T> extends Argument<T> {
+        public static final int INVALID_FILTER = 556;
+        final Argument<T> argument;
+        final Predicate<T> predicate;
+
+        private ArgumentFilter(@NotNull Argument<T> argument, @NotNull Predicate<T> predicate) {
+            super(argument.getId(), argument.allowSpace(), argument.useRemaining());
+            if (argument.getSuggestionCallback() != null)
+                this.setSuggestionCallback(argument.getSuggestionCallback());
+            if (argument.getDefaultValue() != null)
+                this.setDefaultValue(argument.getDefaultValue());
+            this.argument = argument;
+            this.predicate = predicate;
+        }
+
+        @Override
+        public @NotNull T parse(@NotNull String input) throws ArgumentSyntaxException {
+            final T result = argument.parse(input);
+            if (!predicate.test(result))
+                throw new ArgumentSyntaxException("Predicate failed", input, INVALID_FILTER);
+            return result;
+        }
+
+        @Override
+        public void processNodes(@NotNull NodeMaker nodeMaker, boolean executable) {
+            argument.processNodes(nodeMaker, executable);
+        }
+    }
 }

src/main/java/net/minestom/server/command/builder/arguments/ArgumentMap.java

@@ -1,57 +0,0 @@
-package net.minestom.server.command.builder.arguments;
-
-import net.minestom.server.command.builder.NodeMaker;
-import net.minestom.server.command.builder.exception.ArgumentSyntaxException;
-import org.jetbrains.annotations.NotNull;
-
-/**
- * Represents an argument that maps an existing argument to a value.
- *
- * @param <I> The input (any object)
- * @param <O> The output (any object)
- */
-public class ArgumentMap<I, O> extends Argument<O> {
-
-    final Argument<I> argument;
-    final Mapper<I, O> mapper;
-
-    protected ArgumentMap(@NotNull Argument<I> argument, @NotNull Mapper<I, O> mapper) {
-        super(argument.getId(), argument.allowSpace(), argument.useRemaining());
-
-        this.argument = argument;
-        this.mapper = mapper;
-    }
-
-    @Override
-    public @NotNull O parse(@NotNull String input) throws ArgumentSyntaxException {
-        return mapper.accept(argument.parse(input));
-    }
-
-    @Override
-    public void processNodes(@NotNull NodeMaker nodeMaker, boolean executable) {
-        argument.processNodes(nodeMaker, executable);
-    }
-
-    /**
-     * Represents a lambda that can turn an input into an output
-     * that also allows the throwing of ArgumentSyntaxException
-     *
-     * @param <I> The input expected from the Argument
-     * @param <O> The desired output type from this lambda.
-     */
-    @FunctionalInterface
-    public interface Mapper<I, O> {
-
-        /**
-         * Accepts I data from the argument and returns O output
-         *
-         * @param i The input processed from an argument
-         * @return The complex data type that came as a result from this argument
-         * @throws ArgumentSyntaxException If the input can not be turned into the desired output
-         *                                 (E.X. an invalid extension name)
-         */
-        O accept(I i) throws ArgumentSyntaxException;
-
-    }
-
-}

src/test/java/demo/commands/UnloadExtensionCommand.java

@@ -5,7 +5,7 @@ import net.minestom.server.MinecraftServer;
 import net.minestom.server.command.CommandSender;
 import net.minestom.server.command.builder.Command;
 import net.minestom.server.command.builder.CommandContext;
-import net.minestom.server.command.builder.arguments.ArgumentMap;
+import net.minestom.server.command.builder.arguments.Argument;
 import net.minestom.server.command.builder.arguments.ArgumentType;
 import net.minestom.server.command.builder.exception.ArgumentSyntaxException;
 import net.minestom.server.extensions.Extension;
@@ -18,7 +18,7 @@ import java.nio.charset.StandardCharsets;
 
 public class UnloadExtensionCommand extends Command {
 
-    private final ArgumentMap<String, Extension> extensionName;
+    private final Argument<Extension> extensionName;
 
     public UnloadExtensionCommand() {
         super("unload");