[Prototype] Add simpler item function for map

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

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):