CubePool split by color category

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]