diff --git a/exca/map.py b/exca/map.py index a4d337b4..66c757d7 100644 --- a/exca/map.py +++ b/exca/map.py @@ -24,6 +24,7 @@ from . import base, slurm from .cachedict import CacheDict +from .confdict import ConfDict from .dumperloader import DumperLoader MapFunc = tp.Callable[[tp.Sequence[tp.Any]], tp.Iterator[tp.Any]] @@ -305,6 +306,23 @@ def applier(method): return applier + def apply_item( + self, + *, + exclude_from_cache_uid: tp.Iterable[str] | base.ExcludeCallable = (), + cache_type: str | None = None, + ) -> tp.Callable[[C], C]: + params = locals() + params.pop("self") + + def applier(method): + imethod = ItemInfraMethod(method=_UnItemed(method), **params) + self._infra_method = imethod + out = property(imethod) + return out + + return applier + def _find_missing(self, items: tp.Dict[str, tp.Any]) -> tp.Dict[str, tp.Any]: missing = items # deduplicate and check in cache @@ -537,6 +555,49 @@ def check_map_function_output(func: tp.Callable[..., tp.Any]) -> tp.Type[tp.Any] return tp.get_args(annot)[0] # type: ignore +@dataclasses.dataclass +class ItemInfraMethod(base.InfraMethod): + cache_type: str | None = None + params: tp.Tuple[str, ...] = () + + def __post_init__(self) -> None: + super().__post_init__() + if isinstance(self.item_uid, base.Sentinel): + raise ValueError("item_uid needs to be provided") + # self.method = method + + @staticmethod + def item_uid(kwargs: tp.Any) -> str: + return ConfDict(kwargs).to_uid() + + def __call__(self, obj: pydantic.BaseModel) -> tp.Any: + method_override = super().__call__(obj) + return _Itemed(method_override) + + +class _UnItemed: + + def __init__(self, func): + self.func = func + self.__name__ = self.func.__name__ + self.__qualname__ = self.func.__qualname__ + + def __call__(self, obj, items): + for item in items: + yield self.func(obj, **item) + + +class _Itemed: + + def __init__(self, func): + self.func = func + + def __call__(self, *args, **kwargs): + if args: + raise ValueError(f"Only keyword args are supported, got {args=}, {kwargs=}") + return next(iter(self.func([kwargs]))) + + @dataclasses.dataclass class MapInfraMethod(base.InfraMethod): item_uid: tp.Callable[[tp.Any], str] = base.Sentinel() # type: ignore diff --git a/exca/test_map.py b/exca/test_map.py index ed589c78..ea49a474 100644 --- a/exca/test_map.py +++ b/exca/test_map.py @@ -276,6 +276,25 @@ def test_max_workers() -> None: assert isinstance(t_ex._max_workers, int) +class ApplyItem(pydantic.BaseModel): + infra: MapInfra = MapInfra() + param: int = 12 + model_config = pydantic.ConfigDict(extra="forbid") # safer to avoid extra params + + @infra.apply_item(cache_type="MemmapArrayFile") + def process(self, num: int) -> np.ndarray: + print(num) + return np.random.rand(num, self.param) + + +def test_apply_item(tmp_path: Path): + cfg = ApplyItem(infra={"folder": tmp_path}) + out = cfg.process(num=3) + out2 = cfg.process(num=3) + assert out.shape == (3, 12) + np.testing.assert_equal(out, out2) + + def test_missing_item_uid() -> None: # pylint: disable=unused-variable with pytest.raises(TypeError):