Coverage for glotter/download.py: 100%

1from concurrent.futures import ThreadPoolExecutor 

2 

3from glotter.containerfactory import ContainerFactory 

4from glotter.settings import Settings 

5from glotter.source import filter_sources, get_sources 

6 

7 

8def download(args): 

9 def get_key(source): 

10 return f"{source.test_info.container_info.image}:{source.test_info.container_info.tag}" 

11 

12 sources_by_type = filter_sources(args, get_sources(Settings().source_root)) 

13 containers = { 

14 get_key(source): source.test_info.container_info 

15 for sources in sources_by_type.values() 

16 for source in sources 

17 } 

18 if args.parallel: 

19 with ThreadPoolExecutor(max_workers=4) as executor: 

20 executor.map( 

21 lambda source: _download_image_from_source(source, True), 

22 containers.values(), 

23 ) 

24 else: 

25 for container_info in containers.values(): 

26 _download_image_from_source(container_info) 

27 

28 return containers 

29 

30 

31def _download_image_from_source(container_info, parallel=False): 

32 ContainerFactory().get_image(container_info, parallel=parallel) 

33 

34 

35def remove_images(containers, parallel): 

36 if parallel: 

37 with ThreadPoolExecutor(max_workers=4) as executor: 

38 executor.map(ContainerFactory().remove_image, containers.values()) 

39 else: 

40 for container_info in containers.values(): 

41 ContainerFactory().remove_image(container_info)