Limit the type of key values

CommandAPI · JorelAli · Mar 4, 2023 · Jan 29, 2023 · Jan 29, 2023 · Jan 29, 2023
commit 3274f0cfc6e9924e48e957edab56668d7d3cafc8
diff --git a/...kkit/commandapi-bukkit-core/src/main/java/dev/jorel/commandapi/arguments/MapArgument.java b/...kkit/commandapi-bukkit-core/src/main/java/dev/jorel/commandapi/arguments/MapArgument.java
@@ -4,6 +4,7 @@
 import com.mojang.brigadier.arguments.StringArgumentType;
 import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
+import dev.jorel.commandapi.wrappers.MapArgumentKeyType;
 
 import java.util.HashMap;
 import java.util.function.Function;
@@ -26,7 +27,7 @@
 public class MapArgument<K, V> extends Argument<HashMap> implements GreedyArgument {
 
 	private final char delimiter;
-	private final Function<String, K> keyMapper;
+	private final Function<String, ?> keyMapper;
 	private final Function<String, V> valueMapper;
 
 	/**
@@ -35,12 +36,17 @@ public class MapArgument<K, V> extends Argument<HashMap> implements GreedyArgume
 	 * @param nodeName the name to assign to this argument node
 	 * @param delimiter This is used to separate key-value pairs
 	 */
-	MapArgument(String nodeName, char delimiter, Function<String, K> keyMapper, Function<String, V> valueMapper) {
+	MapArgument(String nodeName, char delimiter, MapArgumentKeyType keyType, Function<String, V> valueMapper) {
 		super(nodeName, StringArgumentType.greedyString());
 
 		this.delimiter = delimiter;
-		this.keyMapper = keyMapper;
 		this.valueMapper = valueMapper;
+
+		this.keyMapper = switch (keyType) {
+			case STRING -> (s -> s);
+			case INT -> (Integer::valueOf);
+			case FLOAT -> (Float::valueOf);
+		};
 	}
 
 	@Override
@@ -53,6 +59,7 @@ public CommandAPIArgumentType getArgumentType() {
 		return CommandAPIArgumentType.MAP;
 	}
 
+	@SuppressWarnings("unchecked")
 	@Override
 	public <Source> HashMap<K, V> parseArgument(CommandContext<Source> cmdCtx, String key, Object[] previousArgs) throws CommandSyntaxException {
 		String rawValues = cmdCtx.getArgument(key, String.class);
@@ -78,7 +85,7 @@ public <Source> HashMap<K, V> parseArgument(CommandContext<Source> cmdCtx, Strin
 			visitedCharacters.append(currentChar);
 			if (isAKeyBeingBuilt) {
 				if (currentChar == delimiter) {
-					mapKey = keyMapper.apply(keyBuilder.toString());
+					mapKey = (K) keyMapper.apply(keyBuilder.toString());
 
 					// No need to check the key here because we already know it only consists of letters
 

diff --git a/...mmandapi-bukkit-core/src/main/java/dev/jorel/commandapi/arguments/MapArgumentBuilder.java b/...mmandapi-bukkit-core/src/main/java/dev/jorel/commandapi/arguments/MapArgumentBuilder.java
@@ -1,14 +1,15 @@
 package dev.jorel.commandapi.arguments;
 
+import dev.jorel.commandapi.wrappers.MapArgumentKeyType;
+
 import java.util.function.Function;
 
 /**
  * A builder to create a {@link MapArgument}
  *
- * @param <K>
- * @param <V>
+ * @param <V> The type of values this map should contain
  */
-public class MapArgumentBuilder<K, V> {
+public class MapArgumentBuilder<V> {
 
 	private final String nodeName;
 	private final char delimiter;
@@ -39,19 +40,19 @@ public MapArgumentBuilder(String nodeName, char delimiter) {
 	 *
 	 * @return this map argument builder
 	 */
-	public MapArgumentBuilderValueMapper withKeyMapper(Function<String, K> keyMapper) {
-		return new MapArgumentBuilderValueMapper(keyMapper);
+	public MapArgumentBuilderValueMapper withKeyType(MapArgumentKeyType keyType) {
+		return new MapArgumentBuilderValueMapper(keyType);
 	}
 
 	/**
 	 * An intermediary class for the {@link MapArgumentBuilder}
 	 */
 	public class MapArgumentBuilderValueMapper {
 
-		private final Function<String, K> keyMapper;
+		private final MapArgumentKeyType keyType;
 
-		public MapArgumentBuilderValueMapper(Function<String, K> keyMapper) {
-			this.keyMapper = keyMapper;
+		public MapArgumentBuilderValueMapper(MapArgumentKeyType keyType) {
+			this.keyType = keyType;
 		}
 
 		/**
@@ -81,8 +82,12 @@ public MapArgumentBuilderFinished(Function<String, V> valueMapper) {
 			 *
 			 * @return a new {@link MapArgument}
 			 */
-			public MapArgument<K, V> build() {
-				return new MapArgument<>(nodeName, delimiter, keyMapper, valueMapper);
+			public MapArgument<?, V> build() {
+				return switch (keyType) {
+					case STRING -> new MapArgument<String, V>(nodeName, delimiter, keyType, valueMapper);
+					case FLOAT -> new MapArgument<
EDB7
Float, V>(nodeName, delimiter, keyType, valueMapper);
+					case INT -> new MapArgument<Integer, V>(nodeName, delimiter, keyType, valueMapper);
+				};
 			}
 
 		}

diff --git a/...ommandapi-bukkit-core/src/main/java/dev/jorel/commandapi/wrappers/MapArgumentKeyType.java b/...ommandapi-bukkit-core/src/main/java/dev/jorel/commandapi/wrappers/MapArgumentKeyType.java
@@ -0,0 +1,7 @@
+package dev.jorel.commandapi.wrappers;
+
+public enum MapArgumentKeyType {
+	STRING,
+	FLOAT,
+	INT
+}
diff --git a/...bukkit-test-tests/src/test/java/dev/jorel/commandapi/test/arguments/ArgumentMapTests.java b/...bukkit-test-tests/src/test/java/dev/jorel/commandapi/test/arguments/ArgumentMapTests.java
@@ -8,6 +8,7 @@
 import dev.jorel.commandapi.exceptions.GreedyArgumentException;
 import dev.jorel.commandapi.test.Mut;
 import dev.jorel.commandapi.test.TestBase;
+import dev.jorel.commandapi.wrappers.MapArgumentKeyType;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -46,8 +47,8 @@ public void executionTestWithMapArgument() {
 		Mut<HashMap<String, String>> results = Mut.of();
 
 		new CommandAPICommand("test")
-			.withArguments(new MapArgumentBuilder<String, String>("map", ':')
-				.withKeyMapper(s -> s)
+			.withArguments(new MapArgumentBuilder<String>("map", ':')
+				.withKeyType(MapArgumentKeyType.STRING)
 				.withValueMapper(s -> s)
 				.build()
 			)
@@ -82,8 +83,8 @@ public void executionTestWithMapArgumentAndSpecialValues() {
 		Mut<HashMap<String, String>> results = Mut.of();
 
 		new CommandAPICommand("test")
-			.withArguments(new MapArgumentBuilder<String, String>("map", ':')
-				.withKeyMapper(s -> s)
+			.withArguments(new MapArgumentBuilder<String>("map", ':')
+				.withKeyType(MapArgumentKeyType.STRING)
 				.withValueMapper(s -> s)
 				.build()
 			)
@@ -113,8 +114,8 @@ public void exceptionTestWithMapArgument() {
 		// A GreedyArgumentException should be thrown when it is not the last argument.
 		assertThrows(GreedyArgumentException.class, () -> {
 			new CommandAPICommand("test")
-				.withArguments(new MapArgumentBuilder<String, Integer>("map", '=')
-					.withKeyMapper(s -> s)
+				.withArguments(new MapArgumentBuilder<Integer>("map", '=')
+					.withKeyType(MapArgumentKeyType.STRING)
 					.withValueMapper(Integer::valueOf)
 					.build()
 				)
@@ -124,8 +125,8 @@ public void exceptionTestWithMapArgument() {
 		});
 
 		new CommandAPICommand("test")
-			.withArguments(new MapArgumentBuilder<String, String>("map", ':')
-				.withKeyMapper(s -> s)
+			.withArguments(new MapArgumentBuilder<String>("map", ':')
+				.withKeyType(MapArgumentKeyType.STRING)
 				.withValueMapper(s -> s)
 				.build()
 			)
@@ -150,6 +151,46 @@ public void exceptionTestWithMapArgument() {
 		// Test without closing quotation mark
 		// /test map:"test1
 		assertCommandFailsWith(player, "test map:\"test1", "Could not parse command: A value has to end with a quotation mark! at position 10: map:\"test1<--[HERE]");
+
+		// Test without any quotation marks
+		// /test map:5
+		assertCommandFailsWith(player, "test map:5", "Could not parse command: A value has to start with a quotation mark! at position 4: map:<--[HERE]");
+	}
+
+	@Test
+	public void executionTestWithOtherKeyValuePairs() {
+		Mut<HashMap<String, Integer>> results = Mut.of();
+
+		new CommandAPICommand("test")
+			.withArguments(new MapArgumentBuilder<Integer>("map")
+				.withKeyType(MapArgumentKeyType.STRING)
+				.withValueMapper(Integer::valueOf)
+				.build()
+			)
+			.executesPlayer((player, args) -> {
+				results.set((HashMap<String, Integer>) args.get("map"));
+			})
+			.register();
+
+		PlayerMock player = server.addPlayer();
+
+		// /test map:"598"
+		server.dispatchCommand(player, "test map:\"598\"");
+		Map<String, Integer> testMap = new HashMap<>();
+		testMap.put("map", 598);
+		assertEquals(testMap, results.get());
+
+		// /test map:"598" age:"18"
+		server.dispatchCommand(player, "test map:\"598\" age:\"18\"");
+		testMap.put("age", 18);
+		assertEquals(testMap, results.get());
+
+		// /test map:"598" age:"18" someThirdValue:"19999"
+		server.dispatchCommand(player, "test map:\"598\" age:\"18\" someThirdValue:\"19999\"");
+		testMap.put("someThirdValue", 19999);
+		assertEquals(testMap, results.get());
+
+		assertNoMoreResults(results);
 	}
 
 }