minecraft-kajetans-mod/src/main/java/me/km/snuviscript/commands/LimitCommands.java

package me.km.snuviscript.commands;

import me.hammerle.snuviscript.code.ScriptManager;
import me.km.snuviscript.Scripts;
import net.minecraft.entity.EntityType;

public class LimitCommands {
    public static void registerFunctions(ScriptManager sm, Scripts scripts) {
        sm.registerConsumer("limit.clear", (sc, in) -> {
            scripts.getEntityLimits().clearLimits();
        });
        sm.registerConsumer("limit.set", (sc, in) -> {
            scripts.getEntityLimits().setLimit(getType(in[0].getString(sc)), in[1].getInt(sc));
        });
        sm.registerConsumer("limit.remove", (sc, in) -> {
            scripts.getEntityLimits().removeLimit(getType(in[0].getString(sc)));
        });
    }

    private static EntityType<?> getType(String type) {
        return EntityType.byKey(type).get();
    }
}