spells-mtg 0.10.8__tar.gz → 0.10.9__tar.gz

{spells_mtg-0.10.8 → spells_mtg-0.10.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: spells-mtg
-Version: 0.10.8
+Version: 0.10.9
 Summary: analaysis of 17Lands.com public datasets
 Author-Email: Joel Barnes <oelarnes@gmail.com>
 License: MIT

{spells_mtg-0.10.8 → spells_mtg-0.10.9}/pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
 ]
 requires-python = ">=3.11"
 readme = "README.md"
-version = "0.10.8"
+version = "0.10.9"
 
 [project.license]
 text = "MIT"

{spells_mtg-0.10.8 → spells_mtg-0.10.9}/spells/cache.py

@@ -10,10 +10,17 @@ Caches are cleared per-set when new files are downloaded.
 from enum import StrEnum
 import os
 import sys
-
 import polars as pl
 
 
+class Env(StrEnum):
+    PROD = "prod"
+    TEST = "test"
+
+
+env = Env.PROD
+
+
 class EventType(StrEnum):
     PREMIER = "PremierDraft"
     TRADITIONAL = "TradDraft"
@@ -28,8 +35,30 @@ def spells_print(mode, content):
     print(f"  🪄 {mode} ✨ {content}")
 
 
+def set_test_env():
+    global env
+    env = Env.TEST
+
+
+def set_prod_env():
+    global env
+    env = Env.PROD
+
+
 def data_home() -> str:
     is_win = sys.platform == "win32"
+    global env
+
+    if env == Env.TEST:
+        return os.path.expanduser(
+            os.environ.get(
+                "SPELLS_TEST_HOME",
+                r"~\AppData\Local\SpellsTest"
+                if is_win
+                else "~/.local/share/spellstest/",
+            )
+        )
+
     return os.path.expanduser(
         os.environ.get(
             "SPELLS_DATA_HOME",
@@ -41,6 +70,46 @@ def data_home() -> str:
     )
 
 
+def create_test_data(set_code: str, test_num_drafts: int = 100):
+    """
+    run from prod environment to write test data for `set_code` into
+    the test environment. Then set `SPELLS_DATA_HOME=test_data_home`
+    to run from the test environment
+    """
+
+    context_df = pl.scan_parquet(data_file_path(set_code, "context")).collect()
+    picks_per_pack = context_df["picks_per_pack"][0]
+
+    draft_df = (
+        pl.scan_parquet(data_file_path(set_code, "draft"))
+        .head(50 * (test_num_drafts + 2))
+        .collect()
+    )
+
+    sample_draft_ids = (
+        draft_df.group_by("draft_id")
+        .len()
+        .filter(pl.col("len") == picks_per_pack * 3)["draft_id"][0:test_num_drafts]
+    )
+
+    draft_sample_df = draft_df.filter(pl.col("draft_id").is_in(sample_draft_ids))
+    game_sample_df = (
+        pl.scan_parquet(data_file_path(set_code, "game"))
+        .filter(pl.col("draft_id").is_in(sample_draft_ids))
+        .collect()
+    )
+    card_df = pl.scan_parquet(data_file_path(set_code, "card")).collect()
+
+    set_test_env()
+    if not os.path.isdir(set_dir := external_set_path(set_code)):
+        os.makedirs(set_dir)
+    context_df.write_parquet(data_file_path(set_code, "context"))
+    draft_sample_df.write_parquet(data_file_path(set_code, "draft"))
+    game_sample_df.write_parquet(data_file_path(set_code, "game"))
+    card_df.write_parquet(data_file_path(set_code, "card"))
+    set_prod_env()
+
+
 def data_dir_path(cache_dir: DataDir) -> str:
     """
     Where 17Lands data is stored. MDU_DATA_DIR environment variable is used, if it exists,

{spells_mtg-0.10.8 → spells_mtg-0.10.9}/spells/enums.py

@@ -60,7 +60,7 @@ class ColName(StrEnum):
     PACK_NUM = "pack_num"  # pack_number plus 1
     PICK_NUMBER = "pick_number"
     PICK_NUM = "pick_num"  # pick_number plus 1
-    PICK_INDEX = "pick_index" # 0 - 3 * picks_per_pack - 1
+    PICK_INDEX = "pick_index"  # 0 - 3 * picks_per_pack - 1
     TAKEN_AT = "taken_at"
     NUM_TAKEN = "num_taken"
     NUM_DRAFTS = "num_drafts"