Properly constrain allowed enchantments in /cenchant and /cfish

Author: Earthcomputer

src/main/java/net/earthcomputer/clientcommands/command/CEnchantCommand.java

@@ -8,7 +8,10 @@
 import net.earthcomputer.clientcommands.features.PlayerRandCracker;
 import net.fabricmc.fabric.api.client.command.v2.FabricClientCommandSource;
 import net.minecraft.command.CommandException;
+import net.minecraft.enchantment.Enchantment;
 import net.minecraft.enchantment.EnchantmentLevelEntry;
+import net.minecraft.item.Item;
+import net.minecraft.item.Items;
 import net.minecraft.text.MutableText;
 import net.minecraft.text.Text;
 import net.minecraft.util.Formatting;
@@ -26,10 +29,14 @@ public static void register(CommandDispatcher<FabricClientCommandSource> dispatc
         dispatcher.register(literal("cenchant")
                 .then(literal("--simulate")
                         .redirect(cenchant, ctx -> withFlags(ctx.getSource(), FLAG_SIMULATE, true)))
-                .then(argument("itemAndEnchantmentsPredicate", itemAndEnchantmentsPredicate().withEnchantmentPredicate(ench -> !ench.isTreasure()))
+                .then(argument("itemAndEnchantmentsPredicate", itemAndEnchantmentsPredicate().withEnchantmentPredicate(CEnchantCommand::enchantmentPredicate).constrainMaxLevel())
                         .executes(ctx -> cenchant(ctx.getSource(), getItemAndEnchantmentsPredicate(ctx, "itemAndEnchantmentsPredicate")))));
     }
 
+    private static boolean enchantmentPredicate(Item item, Enchantment ench) {
+        return !ench.isTreasure() && ench.isAvailableForRandomSelection() && (item == Items.BOOK || ench.target.isAcceptableItem(item));
+    }
+
     private static int cenchant(FabricClientCommandSource source, ItemAndEnchantmentsPredicate itemAndEnchantmentsPredicate) throws CommandException {
         if (!Configs.getEnchantingPrediction()) {
             Text text = Text.translatable("commands.cenchant.needEnchantingPrediction")

src/main/java/net/earthcomputer/clientcommands/command/FishCommand.java

@@ -47,7 +47,7 @@ public static void register(CommandDispatcher<FabricClientCommandSource> dispatc
                 .then(argument("goal", clientItemPredicate(registryAccess))
                     .executes(ctx -> addGoal(ctx.getSource(), getClientItemPredicate(ctx, "goal")))))
             .then(literal("add-enchanted-goal")
-                .then(argument("goal", withString(itemAndEnchantmentsPredicate().withItemPredicate(ENCHANTABLE_ITEMS::contains)))
+                .then(argument("goal", withString(itemAndEnchantmentsPredicate().withItemPredicate(ENCHANTABLE_ITEMS::contains).withEnchantmentPredicate((item, ench) -> ench.isAvailableForRandomSelection()).constrainMaxLevel()))
                     .executes(ctx -> addEnchantedGoal(ctx.getSource(), getWithString(ctx, "goal", ItemAndEnchantmentsPredicate.class)))))
             .then(literal("remove-goal")
                 .then(argument("index", integer(1))

src/main/java/net/earthcomputer/clientcommands/command/arguments/ItemAndEnchantmentsPredicateArgumentType.java

@@ -25,6 +25,7 @@
 
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.BiPredicate;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 
@@ -37,7 +38,8 @@ public class ItemAndEnchantmentsPredicateArgumentType implements ArgumentType<It
     private static final DynamicCommandExceptionType ID_INVALID_EXCEPTION = new DynamicCommandExceptionType(id -> Text.translatable("argument.item.id.invalid", id));
 
     private Predicate<Item> itemPredicate = item -> true;
-    private Predicate<Enchantment> enchantmentPredicate = ench -> true;
+    private BiPredicate<Item, Enchantment> enchantmentPredicate = (item, ench) -> true;
+    private boolean constrainMaxLevel = false;
 
     private ItemAndEnchantmentsPredicateArgumentType() {
     }
@@ -51,11 +53,16 @@ public ItemAndEnchantmentsPredicateArgumentType withItemPredicate(Predicate<Item
         return this;
     }
 
-    public ItemAndEnchantmentsPredicateArgumentType withEnchantmentPredicate(Predicate<Enchantment> predicate) {
+    public ItemAndEnchantmentsPredicateArgumentType withEnchantmentPredicate(BiPredicate<Item, Enchantment> predicate) {
         this.enchantmentPredicate = predicate;
         return this;
     }
 
+    public ItemAndEnchantmentsPredicateArgumentType constrainMaxLevel() {
+        this.constrainMaxLevel = true;
+        return this;
+    }
+
     public static ItemAndEnchantmentsPredicate getItemAndEnchantmentsPredicate(CommandContext<?> context, String name) {
         return context.getArgument(name, ItemAndEnchantmentsPredicate.class);
     }
@@ -223,7 +230,7 @@ private Option parseWithWithoutExactly() throws CommandSyntaxException {
         private Enchantment parseEnchantment(boolean suggest, Option option, ItemStack stack) throws CommandSyntaxException {
             List<Enchantment> allowedEnchantments = new ArrayList<>();
             nextEnchantment: for (Enchantment ench : Registries.ENCHANTMENT) {
-                boolean skip = (!ench.isAcceptableItem(stack) && stack.getItem() != Items.BOOK) || !enchantmentPredicate.test(ench);
+                boolean skip = (!ench.isAcceptableItem(stack) && stack.getItem() != Items.BOOK) || !enchantmentPredicate.test(stack.getItem(), ench);
                 if (skip) {
                     continue;
                 }
@@ -278,7 +285,7 @@ private Enchantment parseEnchantment(boolean suggest, Option option, ItemStack s
 
         private int parseEnchantmentLevel(boolean suggest, Option option, ItemStack stack, Enchantment enchantment) throws CommandSyntaxException {
             int maxLevel;
-            {
+            if (constrainMaxLevel) {
                 int enchantability = stack.getItem().getEnchantability();
                 int level = 30 + 1 + enchantability / 4 + enchantability / 4;
                 level += Math.round(level * 0.15f);
@@ -287,6 +294,8 @@ private int parseEnchantmentLevel(boolean suggest, Option option, ItemStack stac
                         break;
                     }
                 }
+            } else {
+                maxLevel = enchantment.getMaxLevel();
             }
 
             List<Integer> allowedLevels = new ArrayList<>();