diff --git a/api/src/main/java/com/viaversion/viaversion/api/data/entity/StoredEntityData.java b/api/src/main/java/com/viaversion/viaversion/api/data/entity/StoredEntityData.java index 8f746fc2f..ec96fa3aa 100644 --- a/api/src/main/java/com/viaversion/viaversion/api/data/entity/StoredEntityData.java +++ b/api/src/main/java/com/viaversion/viaversion/api/data/entity/StoredEntityData.java @@ -51,6 +51,15 @@ public interface StoredEntityData { */ @Nullable T get(Class objectClass); + /** + * Removes and returns an object from the storage if present. + * + * @param objectClass class of the object to remove + * @param object type + * @return removed object if present + */ + @Nullable T remove(Class objectClass); + /** * Stores an object based on its class. * diff --git a/common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityImpl.java b/common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityImpl.java index 8b74b881c..cc1cafa9e 100644 --- a/common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityImpl.java +++ b/common/src/main/java/com/viaversion/viaversion/data/entity/StoredEntityImpl.java @@ -42,6 +42,11 @@ public final class StoredEntityImpl implements StoredEntityData { return (T) storedObjects.get(objectClass); } + @Override + public @Nullable T remove(Class objectClass) { + return (T) storedObjects.remove(objectClass); + } + @Override public boolean has(Class objectClass) { return storedObjects.containsKey(objectClass);