Refactor cube_pool.py out of generate_draft_pool.py

WIll reuse for other modes of play

cube_pool.py

@@ -0,0 +1,29 @@
+# Loads disjoint pools from a cube CSV
+
+import csv
+from collections import defaultdict
+from typing import Dict, List, Set
+
+import cubecobra_csv
+
+
+def load_pools(csv_path: str, pool_tags: Set[str]) -> Dict[str, List[str]]:
+    pools = defaultdict(list)
+    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(', '))
+            # 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:
+                print(f'{card} should only be tagged as one in {pool_tags} '
+                      'but is tagged as {pool_tags_for_card}')
+    return pools

generate_draft_pool.py

@@ -1,15 +1,16 @@
 # Generate a draft pool for dr4ft.info
 
-import csv
 import random
 
 import cube_json
-import cubecobra_csv
+import cube_pool
 
 
 OUTPUT_POOL_SIZE = 360
 CARDS_FROM_OCCASIONAL = 48
 CARDS_FROM_CORE = OUTPUT_POOL_SIZE - CARDS_FROM_OCCASIONAL
+
+
 CUBE_NAME = 'TheElegantCube_fetched'
 
 OUTPUT_FILENAME = 'cards_in_draft.txt'
@@ -18,24 +19,10 @@ OUTPUT_FILENAME = 'cards_in_draft.txt'
 def main():
     csv_path = cube_json.request_cube_csv('TheElegantCube_fetched', 'elegant')
 
-    core_pool = []
+    pools = cube_pool.load_pools(csv_path, {'core', 'occasional'})
-    occasional_pool = []
-    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 = line[cubecobra_csv.COLUMNS['Tags']].split(', ')
-            # print(f'{card} {tags}')
-            if 'occasional' in tags:
-                occasional_pool.append(card)
-            elif 'core' in tags:
-                core_pool.append(card)
-            else:
-                print('{card} is tagged as neither "core" nor "occasional"')
 
-    main_pool = random.sample(core_pool, CARDS_FROM_CORE) + random.sample(occasional_pool, CARDS_FROM_OCCASIONAL)
+    main_pool = (random.sample(pools['core'], CARDS_FROM_CORE) +
+                 random.sample(pools['occasional'], CARDS_FROM_OCCASIONAL))
     with open(OUTPUT_FILENAME, 'w') as f:
         f.write('\n'.join(main_pool))