Coverage for glotter/batch.py: 100%

1import argparse 

2import sys 

3 

4from glotter.download import download, remove_images 

5from glotter.settings import Settings 

6from glotter.source import get_sources 

7from glotter.test import test 

8from glotter.utils import error_and_exit 

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 {source.language.lower() for sources in all_sources.values() for source in sources} 

23 ) 

24 num_languages = len(languages) 

25 num_batches = min(num_languages, args.num_batches) 

26 

27 # Determine starting and ending batch 

28 if args.batch is None: 

29 first_batch = 0 

30 last_batch = num_batches 

31 elif args.batch <= num_batches: 

32 first_batch = args.batch - 1 

33 last_batch = args.batch 

34 else: 

35 sys.exit(0) 

36 

37 # For each batch 

38 exit_code = 0 

39 for n in range(first_batch, last_batch): 

40 # Get languages for this batch 

41 batch_args = argparse.Namespace( 

42 source=None, 

43 project=None, 

44 language=set( 

45 languages[ 

46 (n * num_languages // num_batches) : ((n + 1) * num_languages // num_batches) 

47 ] 

48 ), 

49 parallel=args.parallel, 

50 ) 

51 

52 # Download images for this batch 

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

54 containers = download(batch_args) 

55 

56 # Run tests for this batch 

57 try: 

58 _display_batch("Testing", n, num_batches) 

59 test(batch_args) 

60 except SystemExit as e: 

61 exit_code = exit_code or int(e.code) 

62 

63 # If removing images, remove images for this batch 

64 if args.remove: 

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

66 remove_images(containers, args.parallel) 

67 

68 sys.exit(exit_code) 

69 

70 

71def _display_batch(prefix, n, num_batches): 

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