spells-mtg 0.3.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

spells/draft_data.py

@@ -10,15 +10,17 @@ import datetime
 import functools
 import hashlib
 import re
-from typing import Callable, TypeVar
+from inspect import signature
+from typing import Callable, TypeVar, Any
 
 import polars as pl
+from polars.exceptions import ColumnNotFoundError
 
 from spells.external import data_file_path
 import spells.cache
 import spells.filter
 import spells.manifest
-from spells.columns import ColumnDefinition, ColumnSpec
+from spells.columns import ColDef, ColSpec
 from spells.enums import View, ColName, ColType
 
 
@@ -33,7 +35,7 @@ def _cache_key(args) -> str:
 
 
 @functools.lru_cache(maxsize=None)
-def _get_names(set_code: str) -> tuple[str, ...]:
+def _get_names(set_code: str) -> list[str]:
     card_fp = data_file_path(set_code, View.CARD)
     card_view = pl.read_parquet(card_fp)
     card_names_set = frozenset(card_view.get_column("name").to_list())
@@ -43,7 +45,7 @@ def _get_names(set_code: str) -> tuple[str, ...]:
     cols = draft_view.collect_schema().names()
 
     prefix = "pack_card_"
-    names = tuple(col[len(prefix) :] for col in cols if col.startswith(prefix))
+    names = [col[len(prefix) :] for col in cols if col.startswith(prefix)]
     draft_names_set = frozenset(names)
 
     assert (
@@ -52,34 +54,52 @@ def _get_names(set_code: str) -> tuple[str, ...]:
     return names
 
 
-def _hydrate_col_defs(set_code: str, col_spec_map: dict[str, ColumnSpec]):
-    def get_views(spec: ColumnSpec) -> set[View]:
-        if spec.name == ColName.NAME or spec.col_type in (
-            ColType.AGG,
-            ColType.CARD_SUM,
-        ):
-            return set()
-        if spec.col_type == ColType.CARD_ATTR:
-            return {View.CARD}
-        if spec.views is not None:
-            return set(spec.views)
-        assert (
-            spec.dependencies is not None
-        ), f"Col {spec.name} should have dependencies"
-
-        views = functools.reduce(
-            lambda prev, curr: prev.intersection(curr),
-            [get_views(col_spec_map[dep]) for dep in spec.dependencies],
-        )
+def _get_card_context(set_code: str, col_spec_map: dict[str, ColSpec], card_context: pl.DataFrame | dict[str, dict[str, Any]] | None) -> dict[str, dict[str, Any]]:
+    card_attr_specs = {col:spec for col, spec in col_spec_map.items() if spec.col_type == ColType.CARD_ATTR or spec.name == ColName.NAME}
+    col_def_map = _hydrate_col_defs(set_code, card_attr_specs, card_only=True)
 
-        return views
+    columns = list(col_def_map.keys())
+
+    fp = data_file_path(set_code, View.CARD)
+    card_df = pl.read_parquet(fp)
+    select_rows = _view_select(
+        card_df, frozenset(columns), col_def_map, is_agg_view=False
+    ).to_dicts()
+
+    loaded_context = {row[ColName.NAME]: row for row in select_rows}
+
+    if card_context is not None:
+        if isinstance(card_context, pl.DataFrame):
+            try:
+                card_context = {row[ColName.NAME]: row for row in card_context.to_dicts()}
+            except ColumnNotFoundError:
+                raise ValueError("card_context DataFrame must have column 'name'")
+
+        names = list(loaded_context.keys())
+        for name in names:
+            assert name in card_context, f"card_context must include a row for each card name. {name} missing."
+            for col, value in card_context[name].items():
+                loaded_context[name][col] = value
+
+    return loaded_context
+    
+
+def _determine_expression(spec: ColSpec, names: list[str], card_context: dict[str, dict]) -> pl.Expr | tuple[pl.Expr, ...]:
+    def seed_params(expr):
+        params = {}
+
+        sig_params = signature(expr).parameters
+        if 'names' in sig_params:
+            params['names'] = names
+        if 'card_context' in sig_params:
+            params['card_context'] = card_context
+        return params
+
+    if spec.col_type == ColType.NAME_SUM:
+        if spec.expr is not None:
+            assert isinstance(spec.expr, Callable), f"NAME_SUM column {spec.name} must have a callable `expr` accepting a `name` argument"
+            unnamed_exprs = [spec.expr(**{'name': name, **seed_params(spec.expr)}) for name in names]
 
-    names = _get_names(set_code)
-    assert len(names) > 0, "there should be names"
-    hydrated = {}
-    for key, spec in col_spec_map.items():
-        if spec.col_type == ColType.NAME_SUM and spec.exprMap is not None:
-            unnamed_exprs = map(spec.exprMap, names)
             expr = tuple(
                 map(
                     lambda ex, name: ex.alias(f"{spec.name}_{name}"),
@@ -87,14 +107,76 @@ def _hydrate_col_defs(set_code: str, col_spec_map: dict[str, ColumnSpec]):
                     names,
                 )
             )
-        elif spec.expr is not None:
-            expr = spec.expr.alias(spec.name)
-
         else:
-            if spec.col_type == ColType.NAME_SUM:
-                expr = tuple(map(lambda name: pl.col(f"{spec.name}_{name}"), names))
+            expr = tuple(map(lambda name: pl.col(f"{spec.name}_{name}"), names))
+
+    elif spec.expr is not None:
+        if isinstance(spec.expr, Callable):
+            params = seed_params(spec.expr)
+            if spec.col_type == ColType.PICK_SUM and 'name' in signature(spec.expr).parameters:
+                expr = pl.lit(None)
+                for name in names:
+                    name_params = {'name': name, **params}
+                    expr = pl.when(pl.col(ColName.PICK) == name).then(spec.expr(**name_params)).otherwise(expr)
             else:
-                expr = pl.col(spec.name)
+                expr = spec.expr(**params)
+        else:
+            expr = spec.expr
+        expr = expr.alias(spec.name)
+    else:
+        expr = pl.col(spec.name)
+
+    return expr
+
+
+def _infer_dependencies(name: str, expr: pl.Expr | tuple[pl.Expr,...], col_spec_map: dict[str, ColSpec], names: list[str]) -> set[str]:
+    dependencies = set()
+    tricky_ones = set()
+
+    if isinstance(expr, pl.Expr):
+        dep_cols = [c for c in expr.meta.root_names() if c != name]
+        for dep_col in dep_cols:
+            if dep_col in col_spec_map.keys():
+                dependencies.add(dep_col)
+            else: 
+                tricky_ones.add(dep_col)
+    else:
+        for idx, exp in enumerate(expr):
+            pattern = f"_{names[idx]}$"
+            dep_cols = [c for c in exp.meta.root_names() if c != name]
+            for dep_col in dep_cols:
+                if dep_col in col_spec_map.keys():
+                    dependencies.add(dep_col)
+                elif len(split := re.split(pattern, dep_col)) == 2 and split[0] in col_spec_map:
+                    dependencies.add(split[0])
+                else:
+                    tricky_ones.add(dep_col)
+
+    for item in tricky_ones:
+        found = False
+        for n in names:
+            pattern = f"_{n}$"
+            if not found and len(split := re.split(pattern, item)) == 2 and split[0] in col_spec_map:
+                dependencies.add(split[0])
+                found = True
+        assert found, f"Could not locate column spec for root col {item}" 
+
+    return dependencies
+
+
+def _hydrate_col_defs(set_code: str, col_spec_map: dict[str, ColSpec], card_context: pl.DataFrame | dict[str, dict] | None = None, card_only: bool =False):
+    names = _get_names(set_code)
+
+    if card_only:
+        card_context = {}
+    else:
+        card_context = _get_card_context(set_code, col_spec_map, card_context)
+
+    assert len(names) > 0, "there should be names"
+    hydrated = {}
+    for key, spec in col_spec_map.items():
+        expr = _determine_expression(spec, names, card_context)
+        dependencies = _infer_dependencies(key, expr, col_spec_map, names)
 
         try:
             sig_expr = expr if isinstance(expr, pl.Expr) else expr[0]
@@ -107,22 +189,19 @@ def _hydrate_col_defs(set_code: str, col_spec_map: dict[str, ColumnSpec]):
             else:
                 expr_sig = str(datetime.datetime.now)
 
-        dependencies = tuple(spec.dependencies or ())
-        views = get_views(spec)
         signature = str(
             (
                 spec.name,
                 spec.col_type.value,
                 expr_sig,
-                tuple(view.value for view in views),
                 dependencies,
             )
         )
 
-        cdef = ColumnDefinition(
+        cdef = ColDef(
             name=spec.name,
             col_type=spec.col_type,
-            views=views,
+            views=set(spec.views or set()),
             expr=expr,
             dependencies=dependencies,
             signature=signature,
@@ -134,20 +213,15 @@ def _hydrate_col_defs(set_code: str, col_spec_map: dict[str, ColumnSpec]):
 def _view_select(
     df: DF,
     view_cols: frozenset[str],
-    col_def_map: dict[str, ColumnDefinition],
+    col_def_map: dict[str, ColDef],
     is_agg_view: bool,
-    is_card_sum: bool = False,
 ) -> DF:
     base_cols = frozenset()
     cdefs = [col_def_map[c] for c in view_cols]
     select = []
     for cdef in cdefs:
         if is_agg_view:
-            if (
-                cdef.col_type == ColType.AGG
-                or cdef.col_type == ColType.CARD_SUM
-                and is_card_sum
-            ):
+            if cdef.col_type == ColType.AGG:
                 base_cols = base_cols.union(cdef.dependencies)
                 select.append(cdef.expr)
             else:
@@ -164,7 +238,7 @@ def _view_select(
                 select.append(cdef.expr)
 
     if base_cols != view_cols:
-        df = _view_select(df, base_cols, col_def_map, is_agg_view, is_card_sum)
+        df = _view_select(df, base_cols, col_def_map, is_agg_view)
 
     return df.select(select)
 
@@ -274,44 +348,23 @@ def _base_agg_df(
     )
 
 
-def card_df(
-    set_code: str,
-    extensions: list[ColumnSpec] | None = None,
-):
-    col_spec_map = dict(spells.columns.col_spec_map)
-    if extensions is not None:
-        for spec in extensions:
-            col_spec_map[spec.name] = spec
-
-    col_def_map = _hydrate_col_defs(set_code, col_spec_map)
-
-    columns = [ColName.NAME] + [
-        c for c, cdef in col_def_map.items() if cdef.col_type == ColType.CARD_ATTR
-    ]
-    fp = data_file_path(set_code, View.CARD)
-    card_df = pl.read_parquet(fp)
-    select_df = _view_select(
-        card_df, frozenset(columns), col_def_map, is_agg_view=False
-    )
-    return select_df.select(columns)
-
-
 def summon(
     set_code: str,
     columns: list[str] | None = None,
     group_by: list[str] | None = None,
     filter_spec: dict | None = None,
-    extensions: list[ColumnSpec] | None = None,
+    extensions: list[ColSpec] | None = None,
     use_streaming: bool = False,
     read_cache: bool = True,
     write_cache: bool = True,
+    card_context: pl.DataFrame | dict[str, dict] | None = None
 ) -> pl.DataFrame:
     col_spec_map = dict(spells.columns.col_spec_map)
     if extensions is not None:
         for spec in extensions:
             col_spec_map[spec.name] = spec
 
-    col_def_map = _hydrate_col_defs(set_code, col_spec_map)
+    col_def_map = _hydrate_col_defs(set_code, col_spec_map, card_context)
     m = spells.manifest.create(col_def_map, columns, group_by, filter_spec)
 
     calc_fn = functools.partial(_base_agg_df, set_code, m, use_streaming=use_streaming)
@@ -337,12 +390,6 @@ def summon(
         select_df = _view_select(card_df, card_cols, m.col_def_map, is_agg_view=False)
         agg_df = agg_df.join(select_df, on="name", how="outer", coalesce=True)
 
-        if m.card_sum:
-            card_sum_df = _view_select(
-                agg_df, m.card_sum, m.col_def_map, is_agg_view=True, is_card_sum=True
-            )
-            agg_df = pl.concat([agg_df, card_sum_df], how="horizontal")
-
         if ColName.NAME not in m.group_by:
             agg_df = agg_df.group_by(m.group_by).sum()
 

spells/enums.py

@@ -19,7 +19,6 @@ class ColType(StrEnum):
     NAME_SUM = "name_sum"
     AGG = "agg"
     CARD_ATTR = "card_attr"
-    CARD_SUM = "card_sum"
 
 
 class ColName(StrEnum):
@@ -61,6 +60,7 @@ class ColName(StrEnum):
     PICK_NUM = "pick_num"  # pick_number plus 1
     TAKEN_AT = "taken_at"
     NUM_TAKEN = "num_taken"
+    NUM_DRAFTS = "num_drafts"
     PICK = "pick"
     PICK_MAINDECK_RATE = "pick_maindeck_rate"
     PICK_SIDEBOARD_IN_RATE = "pick_sideboard_in_rate"

spells/manifest.py

@@ -3,18 +3,17 @@ from dataclasses import dataclass
 import spells.columns
 import spells.filter
 from spells.enums import View, ColName, ColType
-from spells.columns import ColumnDefinition
+from spells.columns import ColDef
 
 
 @dataclass(frozen=True)
 class Manifest:
     columns: tuple[str, ...]
-    col_def_map: dict[str, ColumnDefinition]
+    col_def_map: dict[str, ColDef]
     base_view_group_by: frozenset[str]
     view_cols: dict[View, frozenset[str]]
     group_by: tuple[str, ...]
     filter: spells.filter.Filter | None
-    card_sum: frozenset[str]
 
     def __post_init__(self):
         # No name filter check
@@ -40,21 +39,13 @@ class Manifest:
             ), f"Invalid groupby {col}!"
 
         for view, cols_for_view in self.view_cols.items():
-            # cols_for_view are actually in view check
             for col in cols_for_view:
-                assert (
-                    view in self.col_def_map[col].views
-                ), f"View cols generated incorrectly, {col} not in view {view}"
                 # game sum cols on in game, and no NAME groupby
                 assert self.col_def_map[col].col_type != ColType.GAME_SUM or (
                     view == View.GAME and ColName.NAME not in self.base_view_group_by
                 ), f"Invalid manifest for GAME_SUM column {col}"
             if view != View.CARD:
                 for col in self.base_view_group_by:
-                    # base_view_groupbys in view check
-                    assert (
-                        col == ColName.NAME or view in self.col_def_map[col].views
-                    ), f"Groupby {col} not in view {view}!"
                     # base_view_groupbys in view_cols for view
                     assert (
                         col == ColName.NAME or col in cols_for_view
@@ -94,8 +85,8 @@ class Manifest:
 
 def _resolve_view_cols(
     col_set: frozenset[str],
-    col_def_map: dict[str, ColumnDefinition],
-) -> tuple[dict[View, frozenset[str]], frozenset[str]]:
+    col_def_map: dict[str, ColDef],
+) -> dict[View, frozenset[str]]:
     """
     For each view ('game', 'draft', and 'card'), return the columns
     that must be present at the aggregation step. 'name' need not be
@@ -104,7 +95,6 @@ def _resolve_view_cols(
     MAX_DEPTH = 1000
     unresolved_cols = col_set
     view_resolution = {}
-    card_sum = frozenset()
 
     iter_num = 0
     while unresolved_cols and iter_num < MAX_DEPTH:
@@ -116,9 +106,9 @@ def _resolve_view_cols(
                 view_resolution[View.DRAFT] = view_resolution.get(
                     View.DRAFT, frozenset()
                 ).union({ColName.PICK})
-            if cdef.col_type == ColType.CARD_SUM:
-                card_sum = card_sum.union({col})
-            if cdef.views:
+            if cdef.col_type == ColType.CARD_ATTR:
+                view_resolution[View.CARD] = view_resolution.get(View.CARD, frozenset()).union({col})
+            elif cdef.views:
                 for view in cdef.views:
                     view_resolution[view] = view_resolution.get(
                         view, frozenset()
@@ -128,18 +118,42 @@ def _resolve_view_cols(
                     raise ValueError(
                         f"Invalid column def: {col} has neither views nor dependencies!"
                     )
-                for dep in cdef.dependencies:
-                    next_cols = next_cols.union({dep})
+                if cdef.col_type != ColType.AGG:
+                    fully_resolved = True
+                    col_views = frozenset({View.GAME, View.DRAFT, View.CARD})
+                    for dep in cdef.dependencies:
+                        dep_views = frozenset()
+                        for view, view_cols in view_resolution.items():
+                            if dep in view_cols:
+                                dep_views = dep_views.union({view})
+                        if not dep_views:
+                            fully_resolved = False
+                            next_cols = next_cols.union({dep})
+                        else:
+                            col_views = col_views.intersection(dep_views)
+                    if fully_resolved:
+                        assert len(col_views), f"Column {col} can't be defined in any views!"
+                        for view in col_views:
+                            if view not in view_resolution:
+                                print(cdef)
+                                assert False, f"Something went wrong with col {col}"
+
+                            view_resolution[view] = view_resolution[view].union({col})
+                    else:
+                        next_cols = next_cols.union({col})
+                else:
+                    for dep in cdef.dependencies:
+                        next_cols = next_cols.union({dep})
         unresolved_cols = next_cols
 
     if iter_num >= MAX_DEPTH:
         raise ValueError("broken dependency chain in column spec, loop probable")
 
-    return view_resolution, card_sum
+    return view_resolution
 
 
 def create(
-    col_def_map: dict[str, ColumnDefinition],
+    col_def_map: dict[str, ColDef],
     columns: list[str] | None = None,
     group_by: list[str] | None = None,
     filter_spec: dict | None = None,
@@ -148,7 +162,7 @@ def create(
     if columns is None:
         cols = tuple(spells.columns.default_columns)
         if ColName.NAME not in gbs:
-            cols = tuple(c for c in cols if c not in [ColName.COLOR, ColName.RARITY])
+            cols = tuple(c for c in cols if col_def_map[c].col_type != ColType.CARD_ATTR)
     else:
         cols = tuple(columns)
 
@@ -159,12 +173,8 @@ def create(
     if m_filter is not None:
         col_set = col_set.union(m_filter.lhs)
 
-    view_cols, card_sum = _resolve_view_cols(col_set, col_def_map)
     base_view_group_by = frozenset()
 
-    if card_sum:
-        base_view_group_by = base_view_group_by.union({ColName.NAME})
-
     for col in gbs:
         cdef = col_def_map[col]
         if cdef.col_type == ColType.GROUP_BY:
@@ -172,14 +182,23 @@ def create(
         elif cdef.col_type == ColType.CARD_ATTR:
             base_view_group_by = base_view_group_by.union({ColName.NAME})
 
+    view_cols = _resolve_view_cols(col_set, col_def_map)
+
     needed_views = frozenset()
-    for view, cols_for_view in view_cols.items():
-        for col in cols_for_view:
-            if col_def_map[col].views == {view}:  # only found in this view
-                needed_views = needed_views.union({view})
+    if View.CARD in view_cols:
+        needed_views = needed_views.union({View.CARD})
 
-    if not needed_views:
-        needed_views = {View.DRAFT}
+    draft_view_cols = view_cols.get(View.DRAFT, frozenset())
+    game_view_cols = view_cols.get(View.GAME, frozenset())
+
+    base_cols = draft_view_cols.union(game_view_cols)
+
+    if base_cols == draft_view_cols:
+        needed_views = needed_views.union({View.DRAFT})
+    elif base_cols == game_view_cols:
+        needed_views = needed_views.union({View.GAME})
+    else:
+        needed_views = needed_views.union({View.GAME, View.DRAFT})
 
     view_cols = {v: view_cols[v] for v in needed_views}
 
@@ -190,5 +209,4 @@ def create(
         view_cols=view_cols,
         group_by=gbs,
         filter=m_filter,
-        card_sum=card_sum,
     )