A MCEdit filter for batch editing enchantments on items

SmartEnchant.py

# Smart Enchant - A MCEdit filter for batch editing enchantments on items
# Modified by Plastix 
# Original Enchant Filter by SethBling

from pymclevel import MCSchematic
from pymclevel import TileEntity
from pymclevel import TAG_Compound
from pymclevel import TAG_Short
from pymclevel import TAG_Byte
from pymclevel import TAG_String
from pymclevel import TAG_Int
from pymclevel import TAG_List
from numpy import zeros

displayName = "Smart Enchant"
Effects = {
	"Protection": 0,
	"Fire Protection": 1,
	"Feather Falling": 2,
	"Blast Protection": 3,
	"Projectile Protection": 4,
	"Respiration": 5,
	"Aqua Affinity": 6,
	"Thorns": 7,
	"Sharpness": 16,
	"Smite": 17,
	"Bane of Arthropods": 18,
	"Knockback": 19,
	"Fire Aspect": 20,
	"Looting": 21,
	"Efficiency": 32,
	"Silk Touch": 33,
	"Unbreaking": 34,
	"Fortune": 35,
	"Power": 48,
	"Punch": 49,
	"Flame": 50,
	"Infinity": 51,
	"Luck of the Sea": 61,
	"Lure": 62,
	"Make it Glow!": -1, # This will apply no enchantments and make the item glow. Will not remove any enchantments
	}
Modes = {
	"Add Enchantent": 0,
	"Remove Enchantment": 1,
}

inputs = (
	("Mode", tuple(Modes.keys())),
	("Effect", tuple(Effects.keys())),
	("Level", (1, -100, 127)),
	("Determines the item that will be enchanted", "label"),
	("Item ID", (1, 1, 3000)),
	("Deletes all enchantments from selected item", "label"),
	("Remove Enchantments", False),
)

def perform(level, box, options):
	mode = Modes[options["Mode"]]
	effect = Effects[options["Effect"]]
	lvl = options["Level"]
	itemID = options["Item ID"]
	delete = options["Remove Enchantments"]
	
	minx = int(box.minx/16)*16
	minz = int(box.minz/16)*16

	for (chunk, slices, point) in level.getChunkSlices(box):
		for tileEntry in chunk.TileEntities:
			px = tileEntry["x"].value
			py = tileEntry["y"].value
			pz = tileEntry["z"].value

			if px < box.minx or px >= box.maxx:
				continue
			if py < box.miny or py >= box.maxy:
				continue
			if pz < box.minz or pz >= box.maxz:
				continue

			if tileEntry["id"].value == "Trap" or tileEntry["id"].value == "Chest":
				for item in tileEntry["Items"]:
					# Only enchant specified item by ID "itemID"
					if item["id"].value == itemID:
						# Remove the items enchantments if delete is true
						if delete:
							del item["tag"]["ench"]

						# Add the enchantment 'ench' tag and it's parent
						if "tag" not in item:
							item["tag"] = TAG_Compound()
						if "ench" not in item["tag"]:
							item["tag"]["ench"] = TAG_List()

						newEffect = TAG_Compound()
						newEffect["id"] = TAG_Short(effect)
						newEffect["lvl"] = TAG_Short(lvl)
						append = True

						# Remove Enchantment
						if mode == 1:
							del_enchantments = []
							for enchantement in item["tag"]["ench"]:
								if enchantement["id"].value == newEffect["id"].value:
									del_enchantments.append(enchantement)
							for enchantement in del_enchantments:
								item["tag"]["ench"].remove(enchantement)

							# If all the items enchantments have been removed, remove the ench list 
							if len(item["tag"]["ench"]) <= 0:
								del item["tag"]["ench"]
						
						# Add Enchantment	
						else:
							replaced = False
							# Make sure to overwrite the enchantment if we're adding the same one
							for enchantement in item["tag"]["ench"]:
								if enchantement["id"].value == newEffect["id"].value:
									enchantement["lvl"] = TAG_Short(lvl)
									replaced = True

							if not replaced:
								# If the selected enchantment is -1, it means we don't want to add any enchantments, only a enchantment glow
								if newEffect["id"].value != -1 and mode != 1:
									item["tag"]["ench"].append(newEffect)

						chunk.dirty = True