/tests/web/multipartform.py

 1
 2import itertools
 3from mimetypes import guess_type
 4from email.generator import _make_boundary
 5
 6
 7class MultiPartForm(dict):
 8
 9    def __init__(self):
10        self.files = []
11        self.boundary = _make_boundary()
12
13    def get_content_type(self):
14        return "multipart/form-data; boundary=%s" % self.boundary
15
16    def add_file(self, fieldname, filename, fd, mimetype=None):
17        body = fd.read()
18        if mimetype is None:
19            mimetype = guess_type(filename)[0] or "application/octet-stream"
20        self.files.append((fieldname, filename, mimetype, body))
21
22    def bytes(self):
23        parts = []
24        part_boundary = bytearray("--%s" % self.boundary, "ascii")
25
26        # Add the form fields
27        parts.extend([
28            part_boundary,
29            bytearray(
30                "Content-Disposition: form-data; name=\"%s\"" % k,
31                "ascii"
32            ),
33            bytes(),
34            v if isinstance(v, bytes) else bytearray(v, "ascii")
35        ] for k, v in list(self.items()))
36
37        # Add the files to upload
38        parts.extend([
39            part_boundary,
40            bytearray(
41                "Content-Disposition: form-data; name=\"%s\"; filename=\"%s\"" % (
42                    fieldname, filename),
43                "ascii"
44            ),
45            bytearray("Content-Type: %s" % content_type, "ascii"),
46            bytearray(),
47            body if isinstance(body, bytes) else bytearray(body, "ascii"),
48        ] for fieldname, filename, content_type, body in self.files)
49
50        # Flatten the list and add closing boundary marker,
51        # then return CR+LF separated data
52        flattened = list(itertools.chain(*parts))
53        flattened.append(bytearray("--%s--" % self.boundary, "ascii"))
54        res = bytearray()
55        for item in flattened:
56            res += item
57            res += bytearray("\r\n", "ascii")
58        return res