Coverage for python/lsst/pipe/tasks/dataFrameActions/_baseDataFrameActions.py: 44%

36 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-08-06 03:30 +0000

1from __future__ import annotations 

2 

3__all__ = ("DataFrameAction",) 

4 

5from lsst.pex.config import Field, ListField 

6from typing import Iterable, Any, Mapping 

7 

8from lsst.pex.config.configurableActions import ConfigurableAction 

9 

10 

11class DataFrameAction(ConfigurableAction): 

12 _actionCache: Mapping[int, Any] 

13 

14 cache = Field(doc="Controls if the results of this action should be cached," 

15 " only works on frozen actions", 

16 dtype=bool, default=False) 

17 cacheArgs = ListField(doc="If cache is True, this is a list of argument keys that will be used to " 

18 "compute the cache key in addition to the DataFrameId", 

19 dtype=str, optional=True) 

20 

21 def __init_subclass__(cls, **kwargs) -> None: 

22 cls._actionCache = {} 

23 

24 def call_wrapper(function): 

25 def inner_wrapper(self, dataFrame, **kwargs): 

26 dfId = id(dataFrame) 

27 extra = [] 

28 for name in (self.cacheArgs or tuple()): 

29 if name not in kwargs: 

30 raise ValueError(f"{name} is not part of call signature and cant be used for " 

31 "caching") 

32 extra.append(kwargs[name]) 

33 extra.append(dfId) 

34 key = tuple(extra) 

35 if self.cache and self._frozen: 

36 # look up to see if the value is in cache already 

37 if result := self._actionCache.get(key): 

38 return result 

39 result = function(self, dataFrame, **kwargs) 

40 if self.cache and self._frozen: 

41 self._actionCache[key] = result 

42 return result 

43 return inner_wrapper 

44 cls.__call__ = call_wrapper(cls.__call__) 

45 super().__init_subclass__(**kwargs) 

46 

47 def __call__(self, dataFrame, **kwargs) -> Iterable[Any]: 

48 """This method should return the result of an action performed on a 

49 dataframe 

50 """ 

51 raise NotImplementedError("This method should be overloaded in a subclass") 

52 

53 @property 

54 def columns(self) -> Iterable[str]: 

55 """This property should return an iterable of columns needed by this action 

56 """ 

57 raise NotImplementedError("This method should be overloaded in a subclass")