blob: 6ed9d737a88d78e034f6d65bb3b47c26a9ef9250 [file] [log] [blame]
#!/usr/bin/python2.7
#
# Copyright 2012 the Melange authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Downloads GCI student forms."""
import argparse
import os
import re
import shutil
# sets up the environment so we need to do this before importing some
# other modules
import interactive # pylint: disable=relative-import
interactive.setup()
interactive.setDjango()
from google.appengine.datastore import datastore_query
from google.appengine.ext import blobstore
from google.appengine.ext import ndb
from melange.models import profile as profile_model
from soc.modules.gci.models import program as program_model
def WriteBlobAsFile(blob, filename):
"""Write the data from blobstore to local disk.
Args:
blob: a google.appengine.ext.blobstore.BlobInfo object.
filename: path to the filename to write to
"""
dst = open(filename, "w")
src = blob.open()
shutil.copyfileobj(src, dst)
def WriteBlobAsFileUnlessExists(blob, filename):
"""Write the data from blobstore to local disk, if it doesn't already exist.
Args:
blob: a google.appengine.ext.blobstore.BlobInfo object.
filename: path to the filename to write to
"""
# TODO(robert): check size of file vs size of blob?
if not os.path.isfile(filename):
WriteBlobAsFile(blob, filename)
def MakeFilename(student, suffix, blob, directory):
"""Create the filename for a blob.
Files are named based on the username of the student, with the
extension stored in blobstore.
The filename will consist of dash (-) separated components:
- legal name of student
- username of student
- an optional suffix
Characters outside of [0-9A-Za-z-] will be replaced with _.
Args:
student: a Student object
suffix: suffix to add to the filename
blob: a google.appengine.ext.blobstore.BlobInfo object.
directory: directory to put file in
Returns:
string containing name of file to write to.
"""
_, ext = os.path.splitext(blob.filename)
outfile = student.legal_name + "-" + student.profile_id
if suffix:
outfile += "-" + suffix
outfile = re.sub("[^a-zA-Z0-9-]", "_", outfile)
outfile += ext
path = os.path.join(directory, outfile)
return path
def BulkDownloadForms(students, key_function, kind, directory):
"""Bulk download the tax form for a student to local disk.
This is similar to DownloadTaxForm, but fetches all the BlobInfo
objects in one call. It is probably mor efficient.
Args:
students: Student objects to download the forms for
key_function: a function that takes one parameter (a student), and returns
the BlobKey to download
kind: what kind of form is this downloading. will be used as a suffix
on the filename.
directory: Where to put the files.
"""
form_blobkeys = [key_function(s) for s in students]
blobs = blobstore.BlobInfo.get(form_blobkeys)
for student, blob in zip(students, blobs):
path = MakeFilename(student, kind, blob, directory)
# pylint: disable=print-statement
print "Downloading form to '%s'..." % path
WriteBlobAsFileUnlessExists(blob, path)
def DownloadStudentForms(program_path, outputdir):
"""Download all the student tax forms for a program to disk.
Args:
program_path: the name of the program. "google/gci2013"
outputdir: directory to write the forms to. will be created if
it doesn't exist.
"""
outputdir = os.path.abspath(outputdir)
if not os.path.exists(outputdir):
os.mkdir(outputdir)
if not os.path.isdir(outputdir):
# pylint: disable=print-statement
print "Could not create output dir: %s" % outputdir
program = program_model.GCIProgram.get_by_key_name(program_path)
if not program:
# pylint: disable=print-statement
print 'Could not find GCI program "%s"' % program_path
return
program_key = ndb.Key.from_old_key(program.key())
query = profile_model.Profile.query(
profile_model.Profile.program == program_key,
profile_model.Profile.is_student == True,
profile_model.Profile.student_data.number_of_completed_tasks > 0)
cursor = datastore_query.Cursor()
more = True
while more:
results, cursor, more = query.fetch_page(100, start_cursor=cursor)
students_with_enrollment_forms = [s for s in results
if s.student_data.enrollment_form]
students_with_consent_forms = [s for s in results
if s.student_data.consent_form]
BulkDownloadForms(students_with_enrollment_forms,
lambda s: s.student_data.enrollment_form,
"enrollment", outputdir)
BulkDownloadForms(students_with_consent_forms,
lambda s: s.student_data.consent_form,
"consent", outputdir)
def main():
args = argparse.ArgumentParser(description="Download Tax Forms")
args.add_argument("--output", "-o", dest="outputdir", default="/tmp/gciforms",
help="write files to target DIR", metavar="DIR")
args.add_argument("--program", "-p", dest="program_name", default="",
help='full program key name (such as "google/gci2013")',
metavar="PROGRAM", required=True)
args.add_argument("app_id",
help="Name of the application. i.e. s~google-melange")
options = args.parse_args()
interactive.setupOAuthRemote(options.app_id)
DownloadStudentForms(options.program_name, options.outputdir)
if __name__ == "__main__":
main()