Improved finalization (again)

Auxilor · Jan 30, 2021 · d23dc10 · d23dc10
1 parent 46c2e3f
commit d23dc10
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/eco-util/src/main/java/com/willfp/eco/util/display/Display.java b/eco-util/src/main/java/com/willfp/eco/util/display/Display.java
@@ -74,6 +74,7 @@ public void registerRevertModule(@NotNull final Function<ItemStack, ItemStack> f
      */
     public ItemStack display(@NotNull final ItemStack itemStack) {
         if (isFinalized(itemStack)) {
+            unfinalize(itemStack);
             return itemStack;
         }
 
@@ -110,12 +111,14 @@ public ItemStack displayAndFinalize(@NotNull final ItemStack itemStack) {
      */
     public ItemStack revert(@NotNull final ItemStack itemStack) {
         if (isFinalized(itemStack)) {
+            unfinalize(itemStack);
             return itemStack;
         }
 
         for (Function<ItemStack, ItemStack> displayFunction : REVERT_FUNCTIONS) {
             displayFunction.apply(itemStack);
         }
+
         return itemStack;
     }
 
@@ -137,6 +140,24 @@ public ItemStack finalize(@NotNull final ItemStack itemStack) {
         return itemStack;
     }
 
+    /**
+     * Unfinalize an ItemStacks.
+     *
+     * @param itemStack The item.
+     * @return The itemstack.
+     */
+    public ItemStack unfinalize(@NotNull final ItemStack itemStack) {
+        Validate.notNull(finalizeKey, "Key cannot be null!");
+        ItemMeta meta = itemStack.getItemMeta();
+        if (meta == null) {
+            return itemStack;
+        }
+        PersistentDataContainer container = meta.getPersistentDataContainer();
+        container.remove(finalizeKey);
+        itemStack.setItemMeta(meta);
+        return itemStack;
+    }
+
     /**
      * If an item is finalized.
      *