refactor to allow adding enchants to arbitrary tags

src/main/java/me/youhavetrouble/enchantio/enchants/BeheadingEnchant.java

@@ -2,13 +2,17 @@ package me.youhavetrouble.enchantio.enchants;
 
 import io.papermc.paper.registry.RegistryKey;
 import io.papermc.paper.registry.data.EnchantmentRegistryEntry;
+import io.papermc.paper.registry.keys.tags.EnchantmentTagKeys;
 import io.papermc.paper.registry.tag.TagKey;
 import io.papermc.paper.tag.TagEntry;
 import net.kyori.adventure.key.Key;
 import net.kyori.adventure.text.Component;
+import org.bukkit.enchantments.Enchantment;
 import org.bukkit.inventory.EquipmentSlotGroup;
 import org.bukkit.inventory.ItemType;
 
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.Set;
 
 @SuppressWarnings("UnstableApiUsage")
@@ -19,9 +23,9 @@ public class BeheadingEnchant implements EnchantioEnchant {
     private final int anvilCost, weight, maxLevel;
     private final EnchantmentRegistryEntry.EnchantmentCost minimumCost;
     private final EnchantmentRegistryEntry.EnchantmentCost maximumCost;
-    private final boolean canGetFromEnchantingTable;
     private final Set<TagEntry<ItemType>> supportedItemTags;
     private final double chanceToDropHeadPerLevel;
+    private final Set<TagKey<Enchantment>> enchantTagKeys = new HashSet<>();
 
     public BeheadingEnchant(
             int anvilCost,
@@ -37,10 +41,13 @@ public class BeheadingEnchant implements EnchantioEnchant {
         this.weight = weight;
         this.minimumCost = minimumCost;
         this.maximumCost = maximumCost;
-        this.canGetFromEnchantingTable = canGetFromEnchantingTable;
         this.supportedItemTags = supportedItemTags;
         this.maxLevel = maxLevel;
         this.chanceToDropHeadPerLevel = chanceToDropHeadPerLevel;
+
+        if (canGetFromEnchantingTable) {
+            enchantTagKeys.add(EnchantmentTagKeys.IN_ENCHANTING_TABLE);
+        }
     }
 
     @Override
@@ -87,11 +94,6 @@ public class BeheadingEnchant implements EnchantioEnchant {
         return Set.of(EquipmentSlotGroup.ANY);
     }
 
-    @Override
-    public boolean canGetFromEnchantingTable() {
-        return canGetFromEnchantingTable;
-    }
-
     @Override
     public TagKey<ItemType> getTagForSupportedItems() {
         return TagKey.create(RegistryKey.ITEM, Key.key("enchantio:beheading_enchantable"));
@@ -102,4 +104,9 @@ public class BeheadingEnchant implements EnchantioEnchant {
         return supportedItemTags;
     }
 
+    @Override
+    public Set<TagKey<Enchantment>> getEnchantTagKeys() {
+        return Collections.unmodifiableSet(enchantTagKeys);
+    }
+
 }