GH-476 Support type wrappers for native JDA option types. (#476)

Rollczi · Nov 2, 2024 · 9934496 · 9934496
1 parent 5d56811
commit 9934496
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 60 deletions.
diff --git a/litecommands-core/src/dev/rollczi/litecommands/requirement/ScheduledRequirementResolver.java b/litecommands-core/src/dev/rollczi/litecommands/requirement/ScheduledRequirementResolver.java
@@ -30,8 +30,6 @@ class ScheduledRequirementResolver<SENDER> {
     private final BindRegistry bindRegistry;
     private final Scheduler scheduler;
 
-    private final BiMap<Class<?>, ArgumentKey, ParserSet<SENDER, ?>> cachedParserSets = new BiHashMap<>();
-
     ScheduledRequirementResolver(ContextRegistry<SENDER> contextRegistry, ParserRegistry<SENDER> parserRegistry, BindRegistry bindRegistry, Scheduler scheduler) {
         this.contextRegistry = contextRegistry;
         this.parserRegistry = parserRegistry;

diff --git a/litecommands-jda/src/dev/rollczi/litecommands/jda/JDACommandTranslator.java b/litecommands-jda/src/dev/rollczi/litecommands/jda/JDACommandTranslator.java
@@ -15,7 +15,9 @@
 import dev.rollczi.litecommands.meta.MetaHolder;
 import dev.rollczi.litecommands.priority.PrioritizedList;
 import dev.rollczi.litecommands.range.Range;
+import dev.rollczi.litecommands.reflect.type.TypeToken;
 import dev.rollczi.litecommands.shared.Preconditions;
+import java.util.function.Function;
 import net.dv8tion.jda.api.Permission;
 import net.dv8tion.jda.api.entities.User;
 import net.dv8tion.jda.api.events.interaction.command.SlashCommandInteractionEvent;
@@ -42,7 +44,8 @@ public class JDACommandTranslator {
     private static final String DESCRIPTION_NO_GENERATED = "no generated description";
 
     private final Map<Class<?>, JDAType<?>> jdaSupportedTypes = new HashMap<>();
-    private final Map<Class<?>, JDATypeOverlay<?>> jdaTypeOverlays = new HashMap<>();
+    private final Map<Class<?>, JDAType<String>> jdaTypeOverlays = new HashMap<>();
+    private final Map<Class<?>, JDATypeWrapper<?>> jdaTypeWrappers = new HashMap<>();
 
     private final ParserRegistry<User> parserRegistry;
 
@@ -51,12 +54,29 @@ public JDACommandTranslator(ParserRegistry<User> parserRegistry) {
     }
 
     public <T> JDACommandTranslator type(Class<T> type, OptionType optionType, JDATypeMapper<T> mapper) {
-        jdaSupportedTypes.put(type, new JDAType<>(type, optionType, mapper));
+        return this.type(type, optionType, (invocation, option) -> mapper.map(option));
+    }
+
+    public <T> JDACommandTranslator type(Class<T> type, OptionType optionType, JDAContextTypeMapper<T> mapper) {
+        jdaSupportedTypes.put(type, new JDAType<>(optionType, mapper));
         return this;
     }
 
     public <T> JDACommandTranslator typeOverlay(Class<T> type, OptionType optionType, JDATypeMapper<String> mapper) {
-        jdaTypeOverlays.put(type, new JDATypeOverlay<>(type, optionType, mapper));
+        jdaTypeOverlays.put(type, new JDAType<>(optionType, mapper));
+        return this;
+    }
+
+    public <T> JDACommandTranslator typeOverlay(Class<T> type, OptionType optionType, JDAContextTypeMapper<String> mapper) {
+        jdaTypeOverlays.put(type, new JDAType<>(optionType, mapper));
+        return this;
+    }
+
+    /**
+     * Wrappers are only used for wrapping types that are supported by Discord API, and we want to skip LiteCommands parsing.
+     */
+    public <T> JDACommandTranslator typeWrapper(Class<T> type, Function<TypeToken<T>, TypeToken<?>> unwrapper, Function<?, T> wrapper) {
+        jdaTypeWrappers.put(type, new JDATypeWrapper<>(unwrapper, wrapper));
         return this;
     }
 
@@ -176,28 +196,35 @@ private <SENDER> CommandExecutor<SENDER> translateExecutor(CommandRoute<SENDER>
             boolean isRequired = isRequired(argument);
 
             Class<?> parsedType = argument.getType().getRawType();
-            if (jdaSupportedTypes.containsKey(parsedType)) {
-                JDAType<?> jdaType = jdaSupportedTypes.get(parsedType);
-                OptionType optionType = jdaType.optionType();
+            JDATypeWrapper<?> wrapper = jdaTypeWrappers.get(parsedType);
 
-                consumer.translate(optionType, jdaType.mapper(), argumentName, description, isRequired, optionType.canSupportChoices());
-                continue;
+            if (wrapper != null) {
+                parsedType = wrapper.unwrapper(argument.getType()).getRawType();
             }
 
-            if (jdaTypeOverlays.containsKey(parsedType)) {
-                JDATypeOverlay<?> jdaTypeOverlay = jdaTypeOverlays.get(parsedType);
-                OptionType optionType = jdaTypeOverlay.optionType();
-
-                consumer.translate(optionType, jdaTypeOverlay.mapper(), argumentName, description, isRequired, optionType.canSupportChoices());
-                continue;
-            }
+            JDAType<?> type = getType(parsedType);
+            JDAContextTypeMapper<?> mapper = wrapper != null
+                ? (invocation, option) -> wrapper.wrap(type.mapper().map(invocation, option))
+                : type.mapper();
 
-            consumer.translate(OptionType.STRING, option -> option.getAsString(), argumentName, description, isRequired, true);
+            consumer.translate(type.optionType(), mapper, argumentName, description, isRequired, type.optionType().canSupportChoices());
         }
 
         return executor;
     }
 
+    private JDAType<?> getType(Class<?> parsedType) {
+        if (jdaSupportedTypes.containsKey(parsedType)) {
+            return jdaSupportedTypes.get(parsedType);
+        }
+
+        if (jdaTypeOverlays.containsKey(parsedType)) {
+            return jdaTypeOverlays.get(parsedType);
+        }
+
+        return new JDAType<>(OptionType.STRING, (invocation, option) -> option.getAsString());
+    }
+
     private <T> boolean isRequired(Argument<T> argument) {
         if (argument.hasDefaultValue()) {
             return false;
@@ -215,7 +242,7 @@ private <T> boolean isRequired(Argument<T> argument) {
     }
 
     private interface TranslateExecutorConsumer {
-        void translate(OptionType optionType, JDATypeMapper<?> mapper, String argName, String description, boolean isRequired, boolean autocomplete);
+        void translate(OptionType optionType, JDAContextTypeMapper<?> mapper, String argName, String description, boolean isRequired, boolean autocomplete);
     }
 
     JDAParseableInput translateArguments(JDALiteCommand command, SlashCommandInteractionEvent interaction) {
@@ -224,7 +251,7 @@ JDAParseableInput translateArguments(JDALiteCommand command, SlashCommandInterac
             .collect(Collectors.toList());
 
         Map<String, OptionMapping> options = interaction.getOptions().stream()
-            .collect(Collectors.toMap(OptionMapping::getName, option -> option));
+            .collect(Collectors.toMap(optionMapping -> optionMapping.getName(), option -> option));
 
         return new JDAParseableInput(routes, options, command);
     }
@@ -235,7 +262,7 @@ JDASuggestionInput translateSuggestions(CommandAutoCompleteInteraction interacti
             .collect(Collectors.toList());
 
         Map<String, OptionMapping> options = interaction.getOptions().stream()
-            .collect(Collectors.toMap(OptionMapping::getName, option -> option));
+            .collect(Collectors.toMap(optionMapping -> optionMapping.getName(), option -> option));
 
         return new JDASuggestionInput(routes, options, interaction.getFocusedOption());
     }
@@ -260,7 +287,7 @@ Invocation<User> translateInvocation(CommandRoute<User> route, Input<?> argument
 
     static final class JDALiteCommand {
         private final SlashCommandData jdaCommandData;
-        private final Map<JDARoute, JDATypeMapper<?>> jdaArgumentTypeMappers = new HashMap<>();
+        private final Map<JDARoute, JDAContextTypeMapper<?>> jdaArgumentTypeMappers = new HashMap<>();
 
         JDALiteCommand(SlashCommandData jdaCommandData) {
             this.jdaCommandData = jdaCommandData;
@@ -270,17 +297,17 @@ SlashCommandData jdaCommandData() {
             return jdaCommandData;
         }
 
-        Object mapArgument(JDARoute jdaRoute, OptionMapping option) {
-            JDATypeMapper<?> typeMapper = jdaArgumentTypeMappers.get(jdaRoute);
+        Object mapArgument(JDARoute jdaRoute, OptionMapping option, Invocation<?> invocation) {
+            JDAContextTypeMapper<?> typeMapper = jdaArgumentTypeMappers.get(jdaRoute);
 
             if (typeMapper == null) {
                 return null;
             }
 
-            return typeMapper.map(option);
+            return typeMapper.map(invocation, option);
         }
 
-        void addTypeMapper(JDARoute route, JDATypeMapper<?> mapper) {
+        void addTypeMapper(JDARoute route, JDAContextTypeMapper<?> mapper) {
             jdaArgumentTypeMappers.put(route, mapper);
         }
     }
@@ -324,13 +351,11 @@ public int hashCode() {
 
     }
 
-    static final class JDAType<T> {
-        private final Class<T> type;
+    static class JDAType<T> {
         private final OptionType optionType;
-        private final JDATypeMapper<T> mapper;
+        private final JDAContextTypeMapper<T> mapper;
 
-        JDAType(Class<T> type, OptionType optionType, JDATypeMapper<T> mapper) {
-            this.type = type;
+        JDAType(OptionType optionType, JDAContextTypeMapper<T> mapper) {
             this.optionType = optionType;
             this.mapper = mapper;
         }
@@ -339,28 +364,26 @@ public OptionType optionType() {
             return optionType;
         }
 
-        public JDATypeMapper<T> mapper() {
+        public JDAContextTypeMapper<T> mapper() {
             return mapper;
         }
     }
 
-    static final class JDATypeOverlay<T> {
-        private final Class<T> type;
-        private final OptionType optionType;
-        private final JDATypeMapper<String> mapper;
+    static final class JDATypeWrapper<T> {
+        private final Function<TypeToken<T>, TypeToken<?>> unwrapper;
+        private final Function<Object, T> wrapper;
 
-        JDATypeOverlay(Class<T> type, OptionType optionType, JDATypeMapper<String> mapper) {
-            this.type = type;
-            this.optionType = optionType;
-            this.mapper = mapper;
+        public JDATypeWrapper(Function<TypeToken<T>, TypeToken<?>> unwrapper, Function<?, T> wrapper) {
+            this.unwrapper = unwrapper;
+            this.wrapper = (Function<Object, T>) wrapper;
         }
 
-        public OptionType optionType() {
-            return optionType;
+        public TypeToken<?> unwrapper(TypeToken<?> type) {
+            return unwrapper.apply((TypeToken<T>) type);
         }
 
-        public JDATypeMapper<String> mapper() {
-            return mapper;
+        private T wrap(Object object) {
+            return wrapper.apply(object);
         }
     }
 

diff --git a/litecommands-jda/src/dev/rollczi/litecommands/jda/JDAContextTypeMapper.java b/litecommands-jda/src/dev/rollczi/litecommands/jda/JDAContextTypeMapper.java
@@ -0,0 +1,10 @@
+package dev.rollczi.litecommands.jda;
+
+import dev.rollczi.litecommands.invocation.Invocation;
+import net.dv8tion.jda.api.interactions.commands.OptionMapping;
+
+public interface JDAContextTypeMapper<T> {
+
+    T map(Invocation<?> invocation, OptionMapping option);
+
+}
diff --git a/litecommands-jda/src/dev/rollczi/litecommands/jda/JDAParseableInput.java b/litecommands-jda/src/dev/rollczi/litecommands/jda/JDAParseableInput.java
@@ -1,6 +1,5 @@
 package dev.rollczi.litecommands.jda;
 
-import dev.rollczi.litecommands.LiteCommandsException;
 import dev.rollczi.litecommands.argument.Argument;
 import dev.rollczi.litecommands.argument.parser.ParseResult;
 import dev.rollczi.litecommands.argument.parser.input.ParseableInputMatcher;
@@ -9,6 +8,7 @@
 import dev.rollczi.litecommands.argument.parser.input.ParseableInput;
 import dev.rollczi.litecommands.invalidusage.InvalidUsage;
 import dev.rollczi.litecommands.invocation.Invocation;
+import dev.rollczi.litecommands.range.Range;
 import dev.rollczi.litecommands.reflect.ReflectUtil;
 import java.util.function.Supplier;
 import net.dv8tion.jda.api.interactions.commands.OptionMapping;
@@ -47,30 +47,31 @@ public <SENDER, T> ParseResult<T> nextArgument(Invocation<SENDER> invocation, Ar
             OptionMapping optionMapping = arguments.get(argument.getName());
 
             if (optionMapping == null) {
-                return ParseResult.failure(InvalidUsage.Cause.MISSING_ARGUMENT);
+                Parser<SENDER, T> parser = parserProvider.get();
+                Range range = parser.getRange(argument);
+                ParseResult<T> defaultResult = argument.getDefaultValue()
+                    .orElseGet(() -> ParseResult.failure(InvalidUsage.Cause.MISSING_ARGUMENT));
+
+                if (range.getMin() == 0) {
+                    return parser.parse(invocation, argument, RawInput.of())
+                        .mapFailure(failure -> defaultResult);
+                }
+
+                return defaultResult;
             }
 
             Class<T> type = argument.getType().getRawType();
-            Object input = command.mapArgument(toRoute(argument.getName()), optionMapping);
+            Object input = command.mapArgument(toRoute(argument.getName()), optionMapping, invocation);
 
             consumedArguments.add(argument.getName());
 
             if (ReflectUtil.instanceOf(input, type)) {
                 return ParseResult.success((T) input);
             }
 
-            try {
-                Parser<SENDER, T> parser = parserProvider.get();
+            Parser<SENDER, T> parser = parserProvider.get();
 
-                return parser.parse(invocation, argument, RawInput.of(optionMapping.getAsString().split(" ")));
-            }
-            catch (IllegalArgumentException exception) {
-                if (input != null) {
-                    throw new LiteCommandsException("Input: " + input + " is not instance of " + type.getSimpleName());
-                }
-
-                throw new LiteCommandsException("Cannot parse argument: " + argument.getName(), exception);
-            }
+            return parser.parse(invocation, argument, RawInput.of(optionMapping.getAsString().split(" ")));
         }
 
         private JDACommandTranslator.JDARoute toRoute(String argumentName) {

diff --git a/litecommands-jda/src/dev/rollczi/litecommands/jda/JDATypeMapper.java b/litecommands-jda/src/dev/rollczi/litecommands/jda/JDATypeMapper.java
@@ -1,7 +1,15 @@
 package dev.rollczi.litecommands.jda;
 
+import dev.rollczi.litecommands.invocation.Invocation;
 import net.dv8tion.jda.api.interactions.commands.OptionMapping;
 
-public interface JDATypeMapper<T> {
+public interface JDATypeMapper<T> extends JDAContextTypeMapper<T> {
+
     T map(OptionMapping option);
+
+    @Override
+    default T map(Invocation<?> event, OptionMapping option) {
+        return map(option);
+    }
+
 }
diff --git a/litecommands-jda/src/dev/rollczi/litecommands/jda/LiteJDAFactory.java b/litecommands-jda/src/dev/rollczi/litecommands/jda/LiteJDAFactory.java
@@ -11,6 +11,8 @@
 import dev.rollczi.litecommands.jda.permission.DiscordPermissionAnnotationProcessor;
 import dev.rollczi.litecommands.jda.visibility.VisibilityAnnotationProcessor;
 import dev.rollczi.litecommands.scope.Scope;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import net.dv8tion.jda.api.JDA;
 import net.dv8tion.jda.api.entities.Guild;
 import net.dv8tion.jda.api.entities.IMentionable;
@@ -74,10 +76,13 @@ private static JDACommandTranslator createTranslator(ParserRegistry<User> wrappe
             .type(Channel.class, OptionType.CHANNEL, option -> option.getAsChannel())
             .type(GuildChannel.class, OptionType.CHANNEL, option -> option.getAsChannel())
             .type(GuildChannelUnion.class, OptionType.CHANNEL, option -> option.getAsChannel())
+            .type(Member.class, OptionType.USER, (option) -> option.getAsMember())
 
             .typeOverlay(Float.class, OptionType.NUMBER, option -> option.getAsString())
             .typeOverlay(float.class, OptionType.NUMBER, option -> option.getAsString())
-            .typeOverlay(Member.class, OptionType.USER, option -> option.getAsString()) // TODO: Add raw member parer
+
+            .typeWrapper(Optional.class, type -> type.getParameterized(), value -> Optional.of(value))
+            .typeWrapper(CompletableFuture.class, type -> type.getParameterized(), value -> CompletableFuture.completedFuture(value))
             ;
     }