From cba1008b170857da85579f2af12ce07a20229d48 Mon Sep 17 00:00:00 2001
From: henriquenakashima <henrique.nakashima@gmail.com>
Date: Mon, 22 Feb 2021 19:43:50 -0500
Subject: [PATCH] Refactor cube_pool.py out of generate_draft_pool.py

WIll reuse for other modes of play
---
 cube_pool.py           | 29 +++++++++++++++++++++++++++++
 generate_draft_pool.py | 25 ++++++-------------------
 2 files changed, 35 insertions(+), 19 deletions(-)
 create mode 100644 cube_pool.py

diff --git a/cube_pool.py b/cube_pool.py
new file mode 100644
index 0000000..caded89
--- /dev/null
+++ b/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
diff --git a/generate_draft_pool.py b/generate_draft_pool.py
index da2067a..f6037c5 100644
--- a/generate_draft_pool.py
+++ b/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 = []
-    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"')
+    pools = cube_pool.load_pools(csv_path, {'core', '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))