Add restrict_letter_placements option

1 files changed, 26 insertions, 1 deletions

@@ -1,4 +1,13 @@
-from BaseClasses import Location
+from enum import Enum
+
+from BaseClasses import Location, Item
+from .items import Lingo2Item
+
+
+class LetterPlacementType(Enum):
+    ANY = 0
+    DISALLOW = 1
+    FORCE = 2
 
 
 class Lingo2Location(Location):
@@ -6,3 +15,19 @@ class Lingo2Location(Location):
 
     port_id: int
     goal: bool
+    letter_placement_type: LetterPlacementType
+
+    def set_up_letter_rule(self, lpt: LetterPlacementType):
+        self.letter_placement_type = lpt
+        self.item_rule = self._l2_item_rule
+
+    def _l2_item_rule(self, item: Item) -> bool:
+        if not isinstance(item, Lingo2Item):
+            return True
+
+        if self.letter_placement_type == LetterPlacementType.FORCE:
+            return item.is_letter
+        elif self.letter_placement_type == LetterPlacementType.DISALLOW:
+            return not item.is_letter
+
+        return True