Coverage for glotter/batch.py: 100%

1import argparse 

2import sys 

3 

4from glotter.source import get_sources 

5from glotter.settings import Settings 

6from glotter.utils import error_and_exit 

7from glotter.download import download, remove_images 

8from glotter.test import test 

9 

10 

11def batch(args): 

12 # Validate arguments 

13 if args.num_batches < 1: 

14 error_and_exit("Number of batches must be at least 1") 

15 

16 if args.batch is not None and (args.batch < 1 or args.batch > args.num_batches): 

17 error_and_exit(f"Batch number must be from 1 to {args.num_batches}") 

18 

19 # Get all of the languages 

20 all_sources = get_sources(Settings().source_root) 

21 languages = sorted( 

22 { 

23 source.language.lower() 

24 for sources in all_sources.values() 

25 for source in sources 

26 } 

27 ) 

28 num_languages = len(languages) 

29 num_batches = min(num_languages, args.num_batches) 

30 

31 # Determine starting and ending batch 

32 if args.batch is None: 

33 first_batch = 0 

34 last_batch = num_batches 

35 elif args.batch <= num_batches: 

36 first_batch = args.batch - 1 

37 last_batch = args.batch 

38 else: 

39 sys.exit(0) 

40 

41 # For each batch 

42 exit_code = 0 

43 for n in range(first_batch, last_batch): 

44 # Get languages for this batch 

45 batch_args = argparse.Namespace( 

46 source=None, 

47 project=None, 

48 language=set( 

49 languages[ 

50 (n * num_languages // num_batches) : ( 

51 (n + 1) * num_languages // num_batches 

52 ) 

53 ] 

54 ), 

55 parallel=args.parallel, 

56 ) 

57 

58 # Download images for this batch 

59 _display_batch("Downloading images", n, num_batches) 

60 containers = download(batch_args) 

61 

62 # Run tests for this batch 

63 try: 

64 _display_batch("Testing", n, num_batches) 

65 test(batch_args) 

66 except SystemExit as e: 

67 exit_code = exit_code or int(e.code) 

68 

69 # If removing images, remove images for this batch 

70 if args.remove: 

71 _display_batch("Removing images", n, num_batches) 

72 remove_images(containers, args.parallel) 

73 

74 sys.exit(exit_code) 

75 

76 

77def _display_batch(prefix, n, num_batches): 

78 print(f"\n*** {prefix} for batch {n + 1} of {num_batches} ***", flush=True)