add tag for defining enchantments that should not be removed

juancarloscp52 · Feb 21, 2023 · 565c880 · 565c880
1 parent f745378
commit 565c880
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/main/java/me/juancarloscp52/entropy/EntropyTags.java b/src/main/java/me/juancarloscp52/entropy/EntropyTags.java
@@ -14,4 +14,5 @@ public class EntropyTags {
     public static final TagKey<Block> VOID_SIGHT_BREAKS = TagKey.of(Registries.BLOCK.getKey(), new Identifier("entropy", "void_sight_breaks"));
     public static final TagKey<Block> SHOWN_DURING_XRAY = TagKey.of(Registries.BLOCK.getKey(), new Identifier("entropy", "shown_during_xray"));
     public static final TagKey<Enchantment> DO_NOT_ENCHANT_WITH = TagKey.of(Registries.ENCHANTMENT.getKey(), new Identifier("entropy", "do_not_enchant_with"));
+    public static final TagKey<Enchantment> DO_NOT_REMOVE = TagKey.of(Registries.ENCHANTMENT.getKey(), new Identifier("entropy", "do_not_remove"));
 }
diff --git a/src/main/java/me/juancarloscp52/entropy/events/db/RemoveEnchantmentsEvent.java b/src/main/java/me/juancarloscp52/entropy/events/db/RemoveEnchantmentsEvent.java
@@ -1,8 +1,11 @@
 package me.juancarloscp52.entropy.events.db;
 
 import me.juancarloscp52.entropy.Entropy;
+import me.juancarloscp52.entropy.EntropyTags;
 import me.juancarloscp52.entropy.events.AbstractInstantEvent;
+import net.minecraft.enchantment.EnchantmentHelper;
 import net.minecraft.item.ItemStack;
+import net.minecraft.registry.Registries;
 
 public class RemoveEnchantmentsEvent extends AbstractInstantEvent {
 
@@ -31,7 +34,9 @@ private void removeEnchant(ItemStack itemStack) {
         if(nbt==null){return;}
         if(Math.random() > 1.0/3){return;} // one chance over 3 to remove an enchantment.
 
-        nbt.remove("Enchantments");
-        itemStack.setNbt(nbt);
+        var enchantments = EnchantmentHelper.fromNbt(itemStack.getEnchantments());
+
+        enchantments.keySet().removeIf(enchantment -> !Registries.ENCHANTMENT.getEntry(enchantment).isIn(EntropyTags.DO_NOT_REMOVE));
+        EnchantmentHelper.set(enchantments, itemStack);
     }
 }