Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/pchain.py: 80%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6''' 

7Helper classes to implement simple processing pipelines with cacheing. 

8''' 

9 

10 

11class Stage(object): 

12 def __init__(self, f): 

13 self._f = f 

14 self._parent = None 

15 self._cache = {} 

16 

17 def __call__(self, *x, **kwargs): 

18 if kwargs.get('nocache', False): 

19 return self.call_nocache(*x) 

20 

21 if x not in self._cache: 

22 if self._parent is not None: 

23 self._cache[x] = self._f(self._parent(*x[:-1]), *x[-1]) 

24 else: 

25 self._cache[x] = self._f(*x[-1]) 

26 

27 return self._cache[x] 

28 

29 def call_nocache(self, *x): 

30 if self._parent is not None: 

31 return self._f(self._parent.call_nocache(*x[:-1]), *x[-1]) 

32 else: 

33 return self._f(*x[-1]) 

34 

35 def clear(self): 

36 self._cache.clear() 

37 

38 

39class Chain(object): 

40 def __init__(self, *stages): 

41 parent = None 

42 self.stages = [] 

43 for stage in stages: 

44 if not isinstance(stage, Stage): 

45 stage = Stage(stage) 

46 

47 stage._parent = parent 

48 parent = stage 

49 self.stages.append(stage) 

50 

51 def clear(self): 

52 for stage in self.stages: 

53 stage.clear() 

54 

55 def __call__(self, *x, **kwargs): 

56 return self.stages[len(x)-1](*x, **kwargs)