CubePool made of Cards instead of strs

cube_pool.py

@@ -1,7 +1,7 @@
 # Loads disjoint pools from a cube CSV
 
-import collections
 import csv
+import dataclasses
 import itertools
 import random
 from collections import defaultdict
@@ -10,31 +10,49 @@ from typing import Dict, List, Set
 import cubecobra_csv
 
 
+@dataclasses.dataclass
+class Card:
+    name: str
+    color_category: str
+    tags: List[str]
+
+    def __post_init__(self):
+        assert self.color_category in 'wubrgchml0', f'{self}'
+
+    def __repr__(self):
+        return self.name
+
+
 class CubePool:
     def __init__(self):
-        self._categories = defaultdict(list)
+        self._categories: List[Card] = defaultdict(list)
 
-    def add_card(self, card_name: str, color_category: str):
+    def add_card(self, card: Card):
-        assert color_category in 'wubrgchml', f'{card_name}: {color_category}'
+        self._categories[card.color_category].append(card)
-        self._categories[color_category].append(card_name)
 
-    def wide_sample(self, n: int) -> List[str]:
+    def wide_sample(self, n: int) -> List[Card]:
         wide_iterator = list(itertools.chain.from_iterable(self._categories.values()))
         sample = random.sample(wide_iterator, n)
         return sample
 
-    def category_sample(self, color_category: str, n: int) -> List[str]:
+    def category_sample(self, color_category: str, n: int) -> List[Card]:
         return random.sample(self._categories[color_category], n)
 
-    def draw_from_category(self, color_category: str, n: int) -> List[str]:
+    def draw_from_category(self, color_category: str, n: int) -> List[Card]:
         drawn = self.category_sample(color_category, n)
         self.remove_cards(drawn, color_category)
         return drawn
 
-    def remove_cards(self, removed: List[str], color_category: str):
+    def remove_cards(self, removed: List[Card], color_category: str):
-        existing_counter = collections.Counter(self._categories[color_category])
+        existing = self._categories[color_category]
-        removed_counter = collections.Counter(removed)
+        self._categories[color_category] = [c for c in existing if not c in removed]
-        self._categories[color_category] = list((existing_counter - removed_counter).elements())
+
+    def fetch(self, card_name: str, color_category: str) -> Card:
+        found: List[Card] = list(filter(lambda c: c.name == card_name, self._categories[color_category]))
+        assert len(found) >= 1
+        card_taken: Card = found[0]
+        self.remove_cards([card_taken], color_category)
+        return card_taken
 
 
 def load_pools(csv_path: str, pool_tags: Set[str]) -> Dict[str, CubePool]:
@@ -44,19 +62,20 @@ def load_pools(csv_path: str, pool_tags: Set[str]) -> Dict[str, CubePool]:
         header_line = next(reader)
         cubecobra_csv.assert_header(header_line)
         for line in reader:
-            card = cubecobra_csv.get_name(line)
+            card_name = cubecobra_csv.get_name(line)
             tags = cubecobra_csv.get_tags(line)
-            # print(f'{card} {tags}')
+            # print(f'{card_name} {tags}')
             pool_tags_for_card = pool_tags.intersection(tags)
             if len(pool_tags_for_card) == 0:
-                print(f'{card} does not have any of the tags: {pool_tags}')
+                print(f'{card_name} does not have any of the tags: {pool_tags}')
             elif len(pool_tags_for_card) > 1:
-                print(f'{card} should only be tagged as one in {pool_tags} '
+                print(f'{card_name} should only be tagged as one in {pool_tags} '
                       'but is tagged as {pool_tags_for_card}')
             else:
                 (pool_tag,) = pool_tags_for_card
                 color_category = cubecobra_csv.get_color_category(line)
-                # print(f'{card} {color_category}')
+                # print(f'{card_name} {color_category}')
-                pools[pool_tag].add_card(card, color_category)
+                card = Card(card_name, color_category, tags)
+                pools[pool_tag].add_card(card)
     return pools
 

generate_draft_pool.py

@@ -23,9 +23,10 @@ def main():
 
     main_pool = (pools['core'].wide_sample(CARDS_FROM_CORE) +
                  pools['occasional'].wide_sample(CARDS_FROM_OCCASIONAL))
+    main_pool_names = [card.name for card in main_pool]
 
     with open(OUTPUT_FILENAME, 'w') as f:
-        f.write('\n'.join(main_pool))
+        f.write('\n'.join(main_pool_names))
 
 
 if __name__ == '__main__':

rogue_duel.py

@@ -2,7 +2,7 @@
 # improving their decks between games.
 
 import random
-from typing import Dict, List, Tuple
+from typing import Dict, List
 
 import cube_json
 import cube_pool
@@ -12,8 +12,8 @@ CUBE_NAME = 'TheElegantCube_fetched'
 
 class RogueDuelPlayerState:
     def __init__(self):
-        self.colors = []
+        self.colors: List[str] = []
-        self.deck = []
+        self.deck: List[cube_pool.Card] = []
 
 
 class RogueDuelState:
@@ -24,7 +24,7 @@ class RogueDuelState:
 
         csv_path = cube_json.request_cube_csv('TheElegantCube_fetched', 'elegant')
         self.pools: Dict[str, cube_pool.CubePool] = cube_pool.load_pools(csv_path, {'core', 'occasional'})
-        self.discard_pile: List[str] = []
+        self.discard_pile: List[cube_pool.Card] = []
 
 
 def main():
@@ -37,7 +37,7 @@ def main():
 
 
 def _decide_colors(state: RogueDuelState):
-    sealed_pools: Dict[str, List[str]] = {}
+    sealed_pools: Dict[str, List[cube_pool.Card]] = {}
     for color_category in 'wubrg':
         sealed_pools[color_category] = (state.pools['core'].draw_from_category(color_category, 6) +
                                         state.pools['occasional'].draw_from_category(color_category, 2))
@@ -74,7 +74,7 @@ def _pick_a_color(colors_available: List[str], prompt: str) -> str:
 
 
 def _pick_artifacts(state: RogueDuelState):
-    artifact_pool: List[str] = (state.pools['core'].draw_from_category('c', 4) +
+    artifact_pool: List[cube_pool.Card] = (state.pools['core'].draw_from_category('c', 4) +
                                            state.pools['occasional'].draw_from_category('c', 1))
     print(artifact_pool)
     card = _pick_a_card(artifact_pool, 'Player 1, choose an artifact:')
@@ -84,9 +84,9 @@ def _pick_artifacts(state: RogueDuelState):
     state.discard_pile.extend(artifact_pool)
 
 
-def _pick_a_card(cards_available: List[str], prompt: str) -> str:
+def _pick_a_card(cards_available: List[cube_pool.Card], prompt: str) -> cube_pool.Card:
     for i, card in enumerate(cards_available, 1):
-        print(f'[{i:>2}] {card}')
+        print(f'[{i:>2}] {card.name}')
     index_chosen = -1
     while index_chosen < 0 or index_chosen >= len(cards_available):
         index_chosen = int(input(prompt)) - 1
@@ -96,11 +96,12 @@ def _pick_a_card(cards_available: List[str], prompt: str) -> str:
 
 
 def _fill_with_lands(state: RogueDuelState):
-    RAINBOW_DUALS = ['Terramorphic Expanse', 'Evolving Wilds']
+    RAINBOW_DUALS = [state.pools['core'].fetch('Terramorphic Expanse', 'l'),
+                     state.pools['core'].fetch('Evolving Wilds', 'l')]
     state.pools['core'].remove_cards(RAINBOW_DUALS, 'l')
     for player_state, rainbow_land in zip(state.player_states, RAINBOW_DUALS):
         for color in player_state.colors:
-            player_state.deck.extend(5 * [_basic_for_color(color)])
+            player_state.deck.extend([_basic_for_color(color) for _ in range(5)])
         player_state.deck.append(rainbow_land)
         # Add two duals
         duals = _fetch_duals(state, player_state.colors)
@@ -109,7 +110,8 @@ def _fill_with_lands(state: RogueDuelState):
 
 
 def _basic_for_color(color: str) -> str:
-    return {'w': 'Plains', 'u': 'Island', 'b': 'Swamp', 'r': 'Mountain', 'g': 'Forest'}[color]
+    name = {'w': 'Plains', 'u': 'Island', 'b': 'Swamp', 'r': 'Mountain', 'g': 'Forest'}[color]
+    return cube_pool.Card(name, '0', color)
 
 
 def _fetch_duals(state: RogueDuelState, colors: List[str]) -> List[str]: