diff --git a/BaseClasses.py b/BaseClasses.py index 9e1a0a0d..377dee7d 100644 --- a/BaseClasses.py +++ b/BaseClasses.py @@ -1013,6 +1013,17 @@ class CollectionState(): return changed + def add_item(self, item: str, player: int, count: int = 1) -> None: + """ + Adds the item to state. + + :param item: The item to be added. + :param player: The player the item is for. + :param count: How many of the item to add. + """ + assert count > 0 + self.prog_items[player][item] += count + def remove(self, item: Item): changed = self.multiworld.worlds[item.player].remove(self, item) if changed: @@ -1021,6 +1032,33 @@ class CollectionState(): self.blocked_connections[item.player] = set() self.stale[item.player] = True + def remove_item(self, item: str, player: int, count: int = 1) -> None: + """ + Removes the item from state. + + :param item: The item to be removed. + :param player: The player the item is for. + :param count: How many of the item to remove. + """ + assert count > 0 + self.prog_items[player][item] -= count + if self.prog_items[player][item] < 1: + del (self.prog_items[player][item]) + + def set_item(self, item: str, player: int, count: int) -> None: + """ + Sets the item in state equal to the provided count. + + :param item: The item to modify. + :param player: The player the item is for. + :param count: How many of the item to now have. + """ + assert count >= 0 + if count == 0: + del (self.prog_items[player][item]) + else: + self.prog_items[player][item] = count + class EntranceType(IntEnum): ONE_WAY = 1 diff --git a/worlds/AutoWorld.py b/worlds/AutoWorld.py index f0004a9f..6ea6c237 100644 --- a/worlds/AutoWorld.py +++ b/worlds/AutoWorld.py @@ -528,7 +528,7 @@ class World(metaclass=AutoWorldRegister): """Called when an item is collected in to state. Useful for things such as progressive items or currency.""" name = self.collect_item(state, item) if name: - state.prog_items[self.player][name] += 1 + state.add_item(name, self.player) return True return False @@ -536,9 +536,7 @@ class World(metaclass=AutoWorldRegister): """Called when an item is removed from to state. Useful for things such as progressive items or currency.""" name = self.collect_item(state, item, True) if name: - state.prog_items[self.player][name] -= 1 - if state.prog_items[self.player][name] < 1: - del (state.prog_items[self.player][name]) + state.remove_item(name, self.player) return True return False diff --git a/worlds/messenger/__init__.py b/worlds/messenger/__init__.py index 4e2c870d..8df59d9b 100644 --- a/worlds/messenger/__init__.py +++ b/worlds/messenger/__init__.py @@ -428,13 +428,13 @@ class MessengerWorld(World): def collect(self, state: "CollectionState", item: "Item") -> bool: change = super().collect(state, item) if change and "Time Shard" in item.name: - state.prog_items[self.player]["Shards"] += int(item.name.strip("Time Shard ()")) + state.add_item("Shards", self.player, int(item.name.strip("Time Shard ()"))) return change def remove(self, state: "CollectionState", item: "Item") -> bool: change = super().remove(state, item) if change and "Time Shard" in item.name: - state.prog_items[self.player]["Shards"] -= int(item.name.strip("Time Shard ()")) + state.remove_item("Shards", self.player, int(item.name.strip("Time Shard ()"))) return change @classmethod