from __future__ import annotations
import dataclasses
import numbers
import re
from dataclasses import dataclass
from typing import Any
import numpy as np
from directupsampling.grid import GridPoint
[docs]
@dataclass(frozen=True)
class IndexPoint(GridPoint):
"""Immutable, hashable, picklable snapshot index."""
seed: int | None = None
step: int | None = None
[docs]
class IndexView:
"""Live view over the index of a SnapshotContainer.
Holds a direct reference to the container's internal dict so it is always
in sync with insertions and deletions without copying.
"""
def __init__(
self,
snapshots: dict[IndexPoint, Any],
index_names: tuple[str, ...],
) -> None:
"""Initialize IndexView."""
self._snapshots = snapshots
self.names = index_names
[docs]
def coerce(self, index: Any, defaults: dict | None = None) -> IndexPoint:
"""Coerce *index* to an IndexPoint, applying *defaults* for None fields.
Returns
-------
IndexPoint
Raises
------
ValueError
If the resulting index fields do not match ``self.names``.
"""
if defaults is None:
defaults = {}
if isinstance(index, IndexPoint):
updates = {k: v for k, v in defaults.items() if getattr(index, k) is None}
if updates:
index = dataclasses.replace(index, **updates)
elif isinstance(index, dict):
merged = {**defaults, **index}
index = IndexPoint(**{k: merged.get(k) for k in self.names})
else:
index = IndexPoint(**dict(zip(self.names, index, strict=True)))
if set(index.todict()) != set(self.names):
msg = f"Given index {index} must correspond to {self.names}."
raise ValueError(msg)
return index
def __contains__(self, index: IndexPoint) -> bool:
"""Return True if *index* is present in the view.
Returns
-------
bool
"""
return index in self._snapshots
def __iter__(self):
return iter(self._snapshots)
def __len__(self) -> int:
return len(self._snapshots)
def __getitem__(self, key: Any) -> IndexPoint | list[IndexPoint]:
keys = list(self._snapshots)
if isinstance(key, (numbers.Integral, slice)):
return keys[key]
if isinstance(key, (list, np.ndarray)):
key = list(key)
if key and isinstance(key[0], (bool, np.bool_)):
return [keys[i] for i, x in enumerate(key) if x]
if key and isinstance(key[0], numbers.Integral):
return [keys[k] for k in key]
return key # list of IndexPoints
msg = f"IndexView indices must be integers, slices, or lists, not {type(key).__name__}" # noqa: E501
raise TypeError(msg)
__hash__ = None # type: ignore[assignment]
def __eq__(self, other: object) -> bool:
return list(self) == list(other)
def __repr__(self) -> str:
return f"IndexView({list(self._snapshots)!r})"
def __getattr__(self, name: str) -> np.ndarray:
# Only called for names not found normally — i.e. index field names.
try:
return np.array([getattr(idx, name) for idx in self._snapshots])
except AttributeError:
raise AttributeError(name) from None
[docs]
def filter(self, expr: str) -> list[IndexPoint]:
"""Return index points where *expr* evaluates to true.
Parameters
----------
expr : str
Expression string, e.g. ``'volume > 15'`` or
``'volume == 12.0 & temperature == 800'``.
Returns
-------
list[IndexPoint]
Matching index points.
"""
py_expr = re.sub(r"(?<![=!<>&|])&(?![&])", " and ", expr)
py_expr = re.sub(r"(?<![|])\|(?![|])", " or ", py_expr)
return [
idx
for idx in self._snapshots
if eval(py_expr, {}, _coerce_index_dict(idx.todict())) # noqa: S307
]
[docs]
def get_grid(self, names: list[str] | None = None) -> list[tuple]:
"""Return sorted unique grid points for *names*.
Parameters
----------
names : list of str, optional
Index field names to extract. Defaults to all fields except
``'seed'`` and ``'step'``.
Returns
-------
list[tuple]
Sorted unique combinations of the requested field values.
"""
if names is None:
names = self.default_grid_names()
grid = {tuple(getattr(idx, n) for n in names): None for idx in self._snapshots}
return sorted(grid)
[docs]
def default_grid_names(self) -> list[str]:
"""Return index field names excluding ``'seed'`` and ``'step'``.
Returns
-------
list[str]
Field names that represent grid axes, in index order.
"""
exclude = {"seed", "step"}
return [n for n in self.names if n not in exclude]
def _coerce_index_dict(d: dict) -> dict:
"""Cast string values that look numeric to int or float, leave others as-is.
Returns
-------
dict
"""
def _coerce(val: Any) -> Any:
if not isinstance(val, str):
return val
try:
return int(val)
except ValueError:
pass
try:
return float(val)
except ValueError:
return val
return {k: _coerce(v) for k, v in d.items()}