Add keyword to specify program IDs for preview image creation

jwql/jwql_monitors/generate_preview_images.py

@@ -438,6 +438,7 @@ def define_options(parser=None, usage=None, conflict_handler='resolve'):
         parser = argparse.ArgumentParser(usage=usage, conflict_handler=conflict_handler)
 
     parser.add_argument('--overwrite', action='store_true', default=None, help='If set, existing preview images will be re-created and overwritten.')
+    parser.add_argument('-p', '--programs', nargs='+', type=int, help='List of program IDs to generate preview images for. If omitted, all programs will be done.')
     return parser
 
 
@@ -540,20 +541,46 @@ def get_base_output_name(filename_dict):
 
 @log_fail
 @log_info
-def generate_preview_images(overwrite):
+def generate_preview_images(overwrite, programs=None):
     """The main function of the ``generate_preview_image`` module.
     See module docstring for further details.
 
     Parameters
     ----------
     overwrite : bool
         If True, any existing preview images and thumbnails are overwritten
+
+    programs : list
+        List of program ID numbers (e.g. 1068, 01220) for which to generate preview
+        images. If None (the default), preview images are generated for all programs.
     """
+    # Get a list of programs to create preview images for. First, generate a list of all
+    # possible programs. We can compare any user inputs to this list, and if there are no
+    # user inputs, then use this entire list.
+    all_programs = [os.path.basename(item) for item in glob.glob(os.path.join(SETTINGS['filesystem'], 'public', 'jw*'))]
+    all_programs.extend([os.path.basename(item) for item in glob.glob(os.path.join(SETTINGS['filesystem'], 'proprietary', 'jw*'))])
+
+    if programs is None:
+        program_list = all_programs
+    else:
+        if not isinstance(programs, list):
+            raise ValueError(f'program_list must be a list. In this call, it is {type(program_list)}')
+        program_list = []
+        for prog in programs:
+            jwprog = f'jw{str(prog).zfill(5)}'
+            if jwprog in all_programs:
+                program_list.append(jwprog)
+            else:
+                logging.info(f'Program {prog} not present in filesystem. Excluding.')
+
+    if len(program_list) > 0:
+        program_list = sorted(program_list, reverse=True)
+    else:
+        no_prog_message = f'Empty list of programs. No preview images to be made.'
+        logging.info(no_prog_message)
+        raise ValueError(no_prog_message)
 
     # Process programs in parallel
-    program_list = [os.path.basename(item) for item in glob.glob(os.path.join(SETTINGS['filesystem'], 'public', 'jw*'))]
-    program_list.extend([os.path.basename(item) for item in glob.glob(os.path.join(SETTINGS['filesystem'], 'proprietary', 'jw*'))])
-    program_list = list(set(program_list))
     pool = multiprocessing.Pool(processes=int(SETTINGS['cores']))
     program_list = [(element, overwrite) for element in program_list]
     results = pool.starmap(process_program, program_list)
@@ -819,7 +846,7 @@ def update_listfile(filename, file_list, filetype):
 
 
 @lock_module
-def protected_code(overwrite):
+def protected_code(overwrite, programs):
     """Protected code ensures only 1 instance of module will run at any given time
 
     Parameters
@@ -830,11 +857,11 @@ def protected_code(overwrite):
     module = os.path.basename(__file__).strip('.py')
     start_time, log_file = initialize_instrument_monitor(module)
 
-    generate_preview_images(overwrite)
+    generate_preview_images(overwrite, programs=programs)
     update_monitor_table(module, start_time, log_file)
 
 
 if __name__ == '__main__':
     parser = define_options()
     args = parser.parse_args()
-    protected_code(args.overwrite)
+    protected_code(args.overwrite, args.programs)