Address mod-printer v2 PR changes

staff/lab/mod-printer

@@ -6,6 +6,10 @@ import subprocess
 import sys
 
 
+class CUPSException(Exception):
+    pass
+
+
 def modify_printer(args):
     """Perform internal call to CUPS CLI to modify class(es)."""
     # Get available CUPS classes
@@ -44,7 +48,7 @@ def list_printers(args):
                                  stdout=subprocess.PIPE, stdin=subprocess.PIPE)
     lpstat_out = lpstat_proc.stdout.decode('utf-8').strip().split('\n')
 
-    printer = args.printer and args.printer.lower()
+    printer_arg = args.printer and args.printer.lower()
 
     # Parse lpstat output to find all available printers
     # This includes classes and printers (i.e. shared and non-shared printers)
@@ -55,29 +59,33 @@ def list_printers(args):
             continue
         printer_name = printer_name_match.group(1)
 
-        printer_arg = (printer and printer in printer_name) or not printer
-        class_arg = (args.classname and args.classname in printer_name) or not args.classname
+        if printer_arg and printer_arg not in printer_name:
+            continue
+
+        if args.classname and args.classname not in printer_name:
+            continue
 
-        if printer_name not in cups_classes and printer_arg and class_arg:
+        if printer_name not in cups_classes:
             all_printers.append(printer_name)
 
     if args.verbose:
         print('All printers found: %s' % all_printers)
         print('CUPS class-printer mapping found: %s' % cups_classes)
 
     # Find differences between all printers and assigned printers
-    diff_printers = all_printers[:]
+    assigned_printers = []
+    matches_args = lambda name: (not args.printer or args.printer in name) \
+        and (not args.classname or args.classname in name)
     for class_name, class_members in cups_classes.items():
-        for printer_name in class_members[:]:
-            printer_arg = (printer and printer in printer_name) or not printer
-            class_arg = (args.classname and args.classname in printer_name) or not args.classname
-            if printer_arg and class_arg:
-                diff_printers.remove(printer_name)
-            if printer and printer not in printer_name:
-                class_members.remove(printer_name)
+        class_members = [name for name in class_members if matches_args(name)]
+        assigned_printers.extend(class_members)
 
         # Print out all printers belonging to class
-        print('Class %s contains printers: %s' % (class_name, ', '.join(class_members)))
+        if len(class_members):
+            print('Class %s contains printers: %s' % (class_name, ', '.join(class_members)))
+        else:
+            print('Class %s contains no printers' % class_name)
+    diff_printers = set(all_printers) - set(assigned_printers)
 
     # Print out all printers not assigned to a class (i.e. out of service)
     if len(diff_printers):
@@ -89,18 +97,15 @@ def list_printers(args):
     if args.jobs:
         jobs_cmd = ['lpstat', '-o']
         # Specify targets to lpstat command based on classname and printer args
-        if args.classname:
-            jobs_cmd.extend([name for name in all_printers if args.classname in name])
-            jobs_cmd.append(args.classname)
-        if printer:
-            jobs_cmd.extend([name for name in all_printers if printer in name])
+        jobs_cmd.extend(cups_classes.keys())
+        jobs_cmd.extend([printer for printer in all_printers if matches_args(printer)])
         if args.verbose:
             print('Printing jobs with command %s' % jobs_cmd)
 
         # Perform internal jobs call and print result
         jobs_out = do_internal_call(jobs_cmd)
         # Remove blank lines/extraneous whitespace
-        jobs_out = [job for job in jobs_out if len(job)]
+        jobs_out = [job for job in jobs_out if job]
         if len(jobs_out):
             print('\nCurrent print jobs:')
             print('\n'.join(jobs_out))
@@ -141,9 +146,8 @@ def parse_cups_classes(lpstat_raw, args):
             if args.verbose:
                 print('Overriding CUPS classes with parameter: %s' % args.classname)
         else:
-            print('ERROR: Specified class %s could not be found in available classes %s'
+            raise CUPSException('ERROR: Specified class %s could not be found in available classes %s'
                   % (args.classname, cups_classes))
-            return 1
     return cups_classes
 
 
@@ -175,8 +179,7 @@ def main():
     elif args.action == 'list':
         return list_printers(args)
     else:
-        print('ERROR: Invalid action passed %s' % args.action)
-        return 1
+        raise CUPSException('ERROR: Invalid action passed %s' % args.action)
 
 
 if __name__ == '__main__':