import sys
from zipfile import ZipFile
from django.template import Context
from django.template.loader import get_template


class TemplateZipFile(ZipFile, object):
    def __init__(self, file, template_root=None, *args, **kwargs):
        self.template_root = template_root
        return super(TemplateZipFile, self).__init__(file, *args, **kwargs)

    def add_template(self, template_name, filename=None, context=None, compress_type=None):
        if compress_type is not None:
            if compress_type != self.compress_type and sys.version_info < (2, 7):
                raise "Python2.7 is required for individual file compression."
        if context is None:
            c = Context({})
        else:
            c = Context(context)

        if self.template_root is not None:
            template_name = self.template_root + template_name

        template = get_template(template_name)
        render = template.render(c)

        if filename is None:
            if self.template_root is not None:
                arc_filename = template_name.split(self.template_root)[1]
            else:
                arc_filename = template_name.split('/')[-1]
        if compress_type is not None:
            self.writestr(arc_filename, render, compress_type)
        else:
            self.writestr(arc_filename, render)