from cStringIO import StringIO
from email.mime.multipart import MIMEMultipart
from email.generator import Generator


class MIMEMultipartFormData(MIMEMultipart):
    def __init__(self, boundary = None, _subparts = None, **kwargs):
        MIMEMultipart.__init__(
          self,
          _subtype = 'form-data',
          boundary = boundary,
          _subparts = _subparts,
          **kwargs
        )

    def attach(self, subpart):
        if 'MIME-Version' in subpart:
            if subpart['MIME-Version'] != self['MIME-Version']:
                raise ValueError('subpart has incompatible MIME-Version')
            # Note: This isn't strictly necessary, but there is no point in
            # including a MIME-Version header in each subpart.
            del subpart['MIME-Version']
        MIMEMultipart.attach(self, subpart)

    def attach_form_data(self, subpart, name):
        name = name.replace('"', '\\"')
        subpart['Content-Disposition'] = 'form-data; name="%s"' % name
        self.attach(subpart)

    def attach_file(self, subpart, filename):
        filename = filename.replace('"', '\\"')
        subpart['Content-Disposition'] = \
          'file; name="files"; filename="%s"' % filename
        self.attach(subpart)

    def get_body(self, trailing_newline = True):
        f = StringIO()
        generator = Generator(f, mangle_from_ = False)
        generator._dispatch(self)
        if trailing_newline:
            f.write('\n')
        return f.getvalue()

    def as_string(self, trailing_newline = True):
        f = StringIO()
        generator = Generator(f, mangle_from_ = False)
        generator.flatten(self)
        if trailing_newline:
            f.write('\n')
        return f.getvalue()