view tests/multipart.py @ 64:bb8d993376aa

* convenience methods for multipart form * include this in test globals
author Jeff Hammel <jhammel@mozilla.com>
date Thu, 01 Mar 2012 18:15:34 -0800
parents af1476a936fc
children
line wrap: on
line source

"""
from
http://www.doughellmann.com/PyMOTW/urllib2/#uploading-files

Very very sad that python2 stdlib doesn't have this
"""

import itertools
import mimetools
import mimetypes
import os
import urllib
import urllib2
from cStringIO import StringIO

class MultiPartForm(object):
    """Accumulate the data to be used when posting a form."""

    def __init__(self):
        self.form_fields = []
        self.files = []
        self.boundary = mimetools.choose_boundary()

    def get_content_type(self):
        return 'multipart/form-data; boundary=%s' % self.boundary

    def add_field(self, name, value):
        """Add a simple field to the form data."""
        self.form_fields.append((name, value))

    def add_file(self, fieldname, path, mimetype=None):
        filename = os.path.basename(path)
        f = file(path)
        self.add_file_obj(fieldname, filename, f, mimetype)
        f.close()

    def add_file_obj(self, fieldname, filename, fileHandle, mimetype=None):
        """Add a file object to be uploaded."""
        body = fileHandle.read()
        if mimetype is None:
            mimetype = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
        self.files.append((fieldname, filename, mimetype, body))

    def post(self, url):
        request = urllib2.Request(url)
        body = str(self)
        request.add_header('Content-type', self.get_content_type())
        request.add_header('Content-length', len(body))
        request.add_data(body)
        return urllib2.urlopen(request)

    def __str__(self):
        """Return a string representing the form data, including attached files."""
        # Build a list of lists, each containing "lines" of the
        # request.  Each part is separated by a boundary string.
        # Once the list is built, return a string where each
        # line is separated by '\r\n'.
        parts = []
        part_boundary = '--' + self.boundary

        # Add the form fields
        parts.extend(
            [ part_boundary,
              'Content-Disposition: form-data; name="%s"' % name,
              '',
              value,
            ]
            for name, value in self.form_fields
            )

        # Add the files to upload
        parts.extend(
            [ part_boundary,
              'Content-Disposition: file; name="%s"; filename="%s"' % \
                 (field_name, filename),
              'Content-Type: %s' % content_type,
              '',
              body,
            ]
            for field_name, filename, content_type, body in self.files
            )

        # Flatten the list and add closing boundary marker,
        # then return CR+LF separated data
        flattened = list(itertools.chain(*parts))
        flattened.append('--' + self.boundary + '--')
        flattened.append('')
        return '\r\n'.join(flattened)