Coverage for src/pygnd/utils.py: 86%

1"""Small shared helpers used across the pygnd package.""" 

2 

3import contextlib 

4import joblib 

5 

6 

7@contextlib.contextmanager 

8def tqdm_joblib(tqdm_object): 

9 """Context manager that patches joblib to report progress into a tqdm bar. 

10 

11 Args: 

12 tqdm_object: a `tqdm` progress bar instance to update as joblib tasks complete. 

13 

14 Yields: 

15 The same `tqdm_object`, updated automatically as each parallel batch finishes. 

16 """ 

17 

18 class TqdmBatchCompletionCallback(joblib.parallel.BatchCompletionCallBack): 

19 def __call__(self, *args, **kwargs): 

20 tqdm_object.update(n=self.batch_size) 

21 return super().__call__(*args, **kwargs) 

22 

23 old_batch_callback = joblib.parallel.BatchCompletionCallBack 

24 joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback 

25 try: 

26 yield tqdm_object 

27 finally: 

28 joblib.parallel.BatchCompletionCallBack = old_batch_callback 

29 tqdm_object.close()