CubePool split by color category

cube_pool.py

@@ -1,29 +1,47 @@
 # Loads disjoint pools from a cube CSV
 
 import csv
+import itertools
+import random
 from collections import defaultdict
 from typing import Dict, List, Set
 
 import cubecobra_csv
 
 
+class CubePool:
+    def __init__(self):
+        self._categories = defaultdict(list)
+
+    def add_card(self, card_name: str, color_category: str):
+        assert color_category in 'wubrgchml', f'{card_name}: {color_category}'
+        self._categories[color_category].append(card_name)
+
+    def wide_sample(self, n: int):
+        wide_iterator = list(itertools.chain.from_iterable(self._categories.values()))
+        return random.sample(wide_iterator, n)
+
+
 def load_pools(csv_path: str, pool_tags: Set[str]) -> Dict[str, List[str]]:
-    pools = defaultdict(list)
+    pools = defaultdict(CubePool)
     with open(csv_path) as f:
         reader = csv.reader(f)
         header_line = next(reader)
         cubecobra_csv.assert_header(header_line)
         for line in reader:
-            card = line[cubecobra_csv.COLUMNS['Name']]
-            tags = set(line[cubecobra_csv.COLUMNS['Tags']].split(', '))
+            card = cubecobra_csv.get_name(line)
+            tags = cubecobra_csv.get_tags(line)
             # print(f'{card} {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}')
-            elif len(pool_tags_for_card) == 1:
-                (pool_tag,) = pool_tags_for_card
-                pools[pool_tag].append(card)
-            else:
+            elif len(pool_tags_for_card) > 1:
                 print(f'{card} 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}')
+                pools[pool_tag].add_card(card, color_category)
     return pools
+

cubecobra_csv.py

@@ -1,5 +1,8 @@
 # TODO: Make this an actual API
 
+from typing import List, Set
+
+
 EXPECTED_HEADER = [
     'Name',
     'CMC',
@@ -25,3 +28,39 @@ COLUMNS = {column_name: i for i, column_name in enumerate(EXPECTED_HEADER)}
 
 def assert_header(header_line):
     assert header_line == EXPECTED_HEADER
+
+
+_COLOR_CATEGORIES = {
+    'w': 'w',
+    'u': 'u',
+    'b': 'b',
+    'r': 'r',
+    'g': 'g',
+    'm': 'm',
+    'h': 'h',
+    'c': 'c',
+    'l': 'l',
+    'White': 'w',
+    'Blue': 'u',
+    'Black': 'b',
+    'Red': 'r',
+    'Green': 'g',
+    'Multicolor': 'm',
+    'Hybrid': 'h',
+    'Colorless': 'c',
+    'Land': 'land',
+}
+
+
+def get_name(line: List) -> str:
+    return line[COLUMNS['Name']]
+
+
+def get_tags(line: List) -> Set[str]:
+    return set(line[COLUMNS['Tags']].split(', '))
+
+
+def get_color_category(line: List) -> str:
+    color_category = line[COLUMNS['Color Category']]
+    assert color_category in _COLOR_CATEGORIES, f'color_category "{color_category}" unknown'
+    return _COLOR_CATEGORIES[color_category]

generate_draft_pool.py

@@ -21,8 +21,9 @@ def main():
 
     pools = cube_pool.load_pools(csv_path, {'core', 'occasional'})
 
-    main_pool = (random.sample(pools['core'], CARDS_FROM_CORE) +
-                 random.sample(pools['occasional'], CARDS_FROM_OCCASIONAL))
+    main_pool = (pools['core'].wide_sample(CARDS_FROM_CORE) +
+                 pools['occasional'].wide_sample(CARDS_FROM_OCCASIONAL))
+
     with open(OUTPUT_FILENAME, 'w') as f:
         f.write('\n'.join(main_pool))