projects.mako.cc - wikipedia-api-cdsw/blob - mwclient/upload.py

   1 import random
   2 from cStringIO import StringIO
   3
   4
   5 class Upload(object):
   6
   7     """
   8     Base class for upload objects. This class should always be subclassed
   9     by upload classes and its constructor always be called.
  10
  11     Upload classes are file like object/iterators that have additional
  12     variables length and content_type.
  13     """
  14
  15     BLOCK_SIZE = 8192
  16
  17     def __init__(self, length, content_type):
  18         self.length = length
  19         self.content_type = content_type
  20
  21     def __iter__(self):
  22         return self
  23
  24     def next(self):
  25         data = self.read(self.BLOCK_SIZE)
  26         if data == '':
  27             raise StopIteration
  28         return data
  29
  30     @staticmethod
  31     def encode(s):
  32         if type(s) is str:
  33             return s
  34         elif type(s) is unicode:
  35             return s.encode('utf-8')
  36         else:
  37             return s
  38
  39
  40 class UploadRawData(Upload):
  41
  42     """
  43     This upload class is simply a wrapper around StringIO
  44     """
  45
  46     def __init__(self, data, content_type='application/x-www-form-urlencoded'):
  47         self.fstr = StringIO(data)
  48         Upload.__init__(self, len(data), content_type)
  49
  50     def read(self, length=-1):
  51         return self.fstr.read(length)
  52
  53
  54 class UploadDict(UploadRawData):
  55
  56     """
  57     This class creates an x-www-form-urlencoded representation of a dict
  58     and then passes it through its parent UploadRawData
  59     """
  60
  61     def __init__(self, data):
  62         postdata = '&'.join('%s=%s' % (self.encode(i), self.encode(data[i])) for i in data)
  63         UploadRawData.__init__(self, postdata)
  64
  65
  66 class UploadFile(Upload):
  67
  68     """
  69     This class accepts a file with information and a postdata dictionary
  70     and creates a multipart/form-data representation from it.
  71     """
  72     STAGE_FILEHEADER = 0
  73     STAGE_FILE = 1
  74     STAGE_POSTDATA = 2
  75     STAGE_FOOTER = 3
  76     STAGE_DONE = 4
  77
  78     def __init__(self, filefield, filename, filelength, file, data):
  79         self.stage = self.STAGE_FILEHEADER
  80         self.boundary = self.generate_boundary()
  81         self.postdata = self.generate_multipart_from_dict(data)
  82         self.footer = '\r\n--%s--\r\n' % self.boundary
  83         self.fileheader = ('--%s\r\n' % self.boundary +
  84                            'Content-Disposition: form-data; name="%s"; filename="%s"\r\n' %
  85                           (self.encode(filefield), self.encode(filename)) +
  86                            'Content-Type: application/octet-stream\r\n\r\n')
  87         self.file = file
  88         self.length_left = filelength
  89         self.str_data = None
  90
  91         Upload.__init__(self, len(self.fileheader) + filelength + len(self.postdata) + len(self.footer) + 2,
  92                         'multipart/form-data; boundary=' + self.boundary)
  93
  94     def read(self, length):
  95         if self.stage == self.STAGE_DONE:
  96             return ''
  97         elif self.stage != self.STAGE_FILE:
  98             if self.str_data is None:
  99                 if self.stage == self.STAGE_FILEHEADER:
 100                     self.str_data = StringIO(self.fileheader)
 101                 elif self.stage == self.STAGE_POSTDATA:
 102                     self.str_data = StringIO(self.postdata)
 103                 elif self.stage == self.STAGE_FOOTER:
 104                     self.str_data = StringIO(self.footer)
 105             data = self.str_data.read(length)
 106         else:
 107             if self.length_left:
 108                 if length > self.length_left:
 109                     length = self.length_left
 110                 data = self.file.read(length)
 111                 self.length_left -= len(data)
 112             else:
 113                 self.stage += 1
 114                 return '\r\n'
 115
 116         if data == '':
 117             self.stage += 1
 118             self.str_data = None
 119             return self.read(length)
 120         return data
 121
 122     @staticmethod
 123     def generate_boundary():
 124         return '----%s----' % ''.join((random.choice(
 125             'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789')
 126             for i in xrange(32)))
 127
 128     def generate_multipart_from_dict(self, data):
 129         postdata = []
 130         for i in data:
 131             postdata.append('--' + self.boundary)
 132             postdata.append('Content-Disposition: form-data; name="%s"' % self.encode(i))
 133             postdata.append('')
 134             postdata.append(self.encode(data[i]))
 135         return '\r\n'.join(postdata)