gistfile1.py

#!/usr/bin/env python

import argparse
import sys
import re
import csv
import tempfile
import shutil
import subprocess
import math
from path import path
from progressbar import ProgressBar

from pprint import pprint



def confirm(prompt=None, resp=False):
    """prompts for yes or no response from the user. Returns True for yes and
    False for no.

    'resp' should be set to the default value assumed by the caller when
    user simply types ENTER.

    >>> confirm(prompt='Create Directory?', resp=True)
    Create Directory? [y]|n: 
    True
    >>> confirm(prompt='Create Directory?', resp=False)
    Create Directory? [n]|y: 
    False
    >>> confirm(prompt='Create Directory?', resp=False)
    Create Directory? [n]|y: y
    True

    ripped from http://code.activestate.com/recipes/541096/
    """
    
    if prompt is None:
        prompt = 'Confirm'

    if resp:
        prompt = '%s [%s]|%s: ' % (prompt, 'y', 'n')
    else:
        prompt = '%s [%s]|%s: ' % (prompt, 'n', 'y')
        
    while True:
        ans = raw_input(prompt)
        if not ans:
            return resp
        if ans not in ['y', 'Y', 'n', 'N']:
            print 'please enter y or n.'
            continue
        if ans == 'y' or ans == 'Y':
            return True
        if ans == 'n' or ans == 'N':
            return False



def my_check_output(*popenargs, **kwargs):
    # ripped from python 2.7
    if 'stdout' in kwargs:
        raise ValueError('stdout argument not allowed, it will be overridden.')
    process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
    output, unused_err = process.communicate()
    retcode = process.poll()
    if retcode:
        cmd = kwargs.get("args")
        if cmd is None:
            cmd = popenargs[0]
        raise subprocess.CalledProcessError(retcode, cmd)
    return output



def parse_lab_string(s):
    """Takes a lab string like "l3" "LAB03" "3" returns the lab number as an integer."""

    s_stripped = re.sub('\D', '', s)
    if not len(s_stripped):
        raise argparse.ArgumentTypeError("Lab '%s' does not contain a number!" % s)

    n = int(s_stripped)
    if n == 0:
        raise argparse.ArgumentTypeError('Lab number must be >= 1!')

    return n



def parse_section_string(s):
    """Takes a section string like "Sec_A" "a" and returns the section letter."""

    # grab the last character as long as it is alpha
    match = re.search('[A-Za-z]$', s)
    
    if not match:
        raise argparse.ArgumentTypeError("Section '%s' does not contain a section letter!" % s)

    return match.group(0).upper()



def parse_printer_string(s):
    try:
        # just run lpstat on the printer to see if it exists
        # if it doesnt, it will exit 1 and raise an exception
        subprocess.check_call(['lpstat', '-p' + s], stderr=subprocess.STDOUT)
        return s
    except subprocess.CalledProcessError:
        raise argparse.ArgumentTypeError("Printer '%s' is invalid according to lpstat!" % s)



def get_roster(section):
    """Returns a list of tuples (username, email, full name) of the students in a specified section."""

    # filename is ~= ~/rosters/Sec_X
    roster_filename = path('~/rosters/Sec_%s' % section).expand()

    with open(roster_filename) as f:
        roster_reader = csv.reader(f, delimiter='\t')

        # copy the list so we can iterate over it multiple times
        rows = list(roster_reader)

        # make sure each row has all the fields (username, email, name)
        for row in rows:
            if len(row) != 3:
                raise ValueError('Not enough columns in roster file %s' % roster_filename)

        return map(tuple, rows)



def get_results_files(lab, username_list):
    d = {}

    for user in username_list:
        results_dir = path('~/results/Lab%02d/%s' % (lab, user)).expand()
        # get a list of the result files
        files = results_dir.walkfiles('*.results')
        # add the files under this user
        d[user] = tuple(files)

    return d



def generate_ps_files(result_files, lab):
    # make a temp directory to store the postscript files
    ps_dir = path(tempfile.mkdtemp('_ece364_results'))

    pbar = ProgressBar()

    virtual_pages = 0
    pages = 0
    duplex_pages = 0

    print
    print "Generating PostScript files from result files..."

    for (user, results) in pbar(result_files.items()):
        out_ps = path('%s/%s.ps' % (ps_dir, user))
        # this will raise an exception w/ a non-zero exit value from a2ps
        output = my_check_output(['a2ps', "--left-title=",
                                          "--left-footer=",
                                          "--footer=$n",
                                          "--header=ECE 364: Fall 2010 Lab \\#%d User: %s" % (lab, user),
                                          "-o%s" % out_ps] + list(results),
                                 stderr=subprocess.STDOUT)
        
        match = re.search('\[Total: (\d+) pages on (\d+) sheets\]', output)
        if not match:
            raise ValueError("Couldn't parse a2ps ouput! Output follows:\n%s" % output)

        virtual_pages += int(match.group(1))
        pages += int(match.group(2))
        duplex_pages += int(math.ceil(int(match.group(2)) / 2.0))

    return ps_dir, virtual_pages, pages, duplex_pages



def print_results(ps_dir, printer):

    pbar = ProgressBar()

    # get a list of all the postscript files
    ps_files = list(ps_dir.walkfiles('*.ps'))

    for f in pbar(ps_files):
        # call lpr to print
        subprocess.check_call(['lpr', '-P%s' % printer,
                                      '-omedia=Letter',
                                      '-oStaple=1PLU',
                                      str(f)],
                              stderr=subprocess.STDOUT)

    print

    return



def main():
    arg_parser = argparse.ArgumentParser(description='Print the grading script results for a lab.')

    arg_parser.add_argument('-l', '--lab', required=True, type=parse_lab_string)
    arg_parser.add_argument('-s', '--section', action='append', type=parse_section_string)
    arg_parser.add_argument('-u', '--user', action='append')
    arg_parser.add_argument('-p', '--printer', required=True, type=parse_printer_string)

    args = arg_parser.parse_args()

    # construct a list of the users for which we want to print results
    username_list = []

    # add all the users from a section
    if args.section != None:
        for s in args.section:
            roster = get_roster(s)
            username_list += [roster_tuple[0] for roster_tuple in roster]

    # add any users specified on the command line
    if args.user != None:
        username_list += args.user

    # find the results files
    result_files = get_results_files(args.lab, username_list)

    # actually generate the postscript files and get the page stats
    ps_dir, virtual_pages, pages, duplex_pages = generate_ps_files(result_files, args.lab)

    print
    print "--- STATS ---"
    print "    # students: %d" % len(username_list)
    #print "    # result files: 
    print "    # virtual pages: %d" % virtual_pages
    print "    # real pages: %d" % pages
    print "    # pages w/ duplex printing: %d" % duplex_pages
    print

    if confirm(prompt="Do you want to print off this batch?", resp=False):
        print "OK! Printing the results..."
        print
        # this will actually PRINT PAGES! be careful!
        print_results(ps_dir, args.printer)
    else:
        print "Well I'm glad I asked! Refusing to print results."
        print
    

    # CAREFUL! this is just like `rm -rf $ps_dir`
    shutil.rmtree(ps_dir)

    return 0



if __name__ == "__main__":
    sys.exit(main())