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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: spells-mtg
-Version: 0.10.7
+Version: 0.10.9
 Summary: analaysis of 17Lands.com public datasets
 Author-Email: Joel Barnes <oelarnes@gmail.com>
 License: MIT
@@ -412,6 +412,7 @@ A table of all included columns. Columns can be referenced by enum or by string
 | `PACK_NUM`               | `"pack_num"`               | `DRAFT` | `GROUP_BY`    | 1-indexed                                    | Int      |
 | `PICK_NUMBER`            | `"pick_number"`            | `DRAFT` | `FILTER_ONLY` | Dataset Column                               | Int      |
 | `PICK_NUM`               | `"pick_num"`               | `DRAFT` | `GROUP_BY`    | 1-indexed                                    | Int      |
+| `PICK_INDEX`              | `"pick_index"`            | `DRAFT` | `GROUP_BY`  | 0-indexed, through 39/42/45 depending on set   | Int      |
 | `TAKEN_AT`               | `"taken_at`                | `DRAFT` | `PICK_SUM`    | Summable alias of `PICK_NUM`                 | Int      |
 | `NUM_DRAFTS`           | `"num_drafts"` | `DRAFT` | `PICK_SUM` | | Int |
 | `NUM_TAKEN`              | `"num_taken"`              | `DRAFT` | `PICK_SUM`    | Sum 1 over rows                              | Int      |

{spells_mtg-0.10.7 → spells_mtg-0.10.9}/README.md

@@ -401,6 +401,7 @@ A table of all included columns. Columns can be referenced by enum or by string
 | `PACK_NUM`               | `"pack_num"`               | `DRAFT` | `GROUP_BY`    | 1-indexed                                    | Int      |
 | `PICK_NUMBER`            | `"pick_number"`            | `DRAFT` | `FILTER_ONLY` | Dataset Column                               | Int      |
 | `PICK_NUM`               | `"pick_num"`               | `DRAFT` | `GROUP_BY`    | 1-indexed                                    | Int      |
+| `PICK_INDEX`              | `"pick_index"`            | `DRAFT` | `GROUP_BY`  | 0-indexed, through 39/42/45 depending on set   | Int      |
 | `TAKEN_AT`               | `"taken_at`                | `DRAFT` | `PICK_SUM`    | Summable alias of `PICK_NUM`                 | Int      |
 | `NUM_DRAFTS`           | `"num_drafts"` | `DRAFT` | `PICK_SUM` | | Int |
 | `NUM_TAKEN`              | `"num_taken"`              | `DRAFT` | `PICK_SUM`    | Sum 1 over rows                              | Int      |

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

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

{spells_mtg-0.10.7 → 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.7 → spells_mtg-0.10.9}/spells/columns.py

@@ -166,6 +166,10 @@ _specs: dict[str, ColSpec] = {
         col_type=ColType.GROUP_BY,
         expr=pl.col(ColName.PICK_NUMBER) + 1,
     ),
+    ColName.PICK_INDEX: ColSpec(
+        col_type=ColType.GROUP_BY,
+        expr=lambda set_context: pl.col(ColName.PICK_NUMBER) + pl.col(ColName.PACK_NUMBER) * set_context['picks_per_pack']
+    ),
     ColName.TAKEN_AT: ColSpec(
         col_type=ColType.PICK_SUM,
         expr=pl.col(ColName.PICK_NUM),

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

@@ -60,6 +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
     TAKEN_AT = "taken_at"
     NUM_TAKEN = "num_taken"
     NUM_DRAFTS = "num_drafts"