103 lines
2.9 KiB
Python
103 lines
2.9 KiB
Python
from __future__ import unicode_literals
|
|
|
|
import zipfile
|
|
|
|
try:
|
|
import zlib
|
|
COMPRESSION = zipfile.ZIP_DEFLATED
|
|
except:
|
|
COMPRESSION = zipfile.ZIP_STORED
|
|
|
|
try:
|
|
from cStringIO import StringIO
|
|
except ImportError:
|
|
from StringIO import StringIO
|
|
|
|
from django.core.files.uploadedfile import SimpleUploadedFile
|
|
|
|
|
|
class NotACompressedFile(Exception):
|
|
pass
|
|
|
|
|
|
class CompressedFile(object):
|
|
def __init__(self, file_input=None):
|
|
if file_input:
|
|
try:
|
|
# Is it a file like object?
|
|
file_input.seek(0)
|
|
except AttributeError:
|
|
# If not, try open it.
|
|
self._open(file_input)
|
|
else:
|
|
self.file_object = file_input
|
|
else:
|
|
self._create()
|
|
|
|
def _create(self):
|
|
self.descriptor = StringIO()
|
|
self.zf = zipfile.ZipFile(self.descriptor, mode='w')
|
|
|
|
def _open(self, file_input):
|
|
try:
|
|
# Is it a file like object?
|
|
file_input.seek(0)
|
|
except AttributeError:
|
|
# If not, try open it.
|
|
self.descriptor = open(file_input, 'r+b')
|
|
else:
|
|
self.descriptor = file_input
|
|
|
|
try:
|
|
test = zipfile.ZipFile(self.descriptor, mode='r')
|
|
except zipfile.BadZipfile:
|
|
raise NotACompressedFile
|
|
else:
|
|
test.close()
|
|
self.descriptor.seek(0)
|
|
self.zf = zipfile.ZipFile(self.descriptor, mode='a')
|
|
|
|
def add_file(self, file_input, arcname=None):
|
|
try:
|
|
# Is it a file like object?
|
|
file_input.seek(0)
|
|
except AttributeError:
|
|
# If not, keep it
|
|
self.zf.write(file_input, arcname=arcname, compress_type=COMPRESSION)
|
|
else:
|
|
self.zf.writestr(arcname, file_input.read())
|
|
|
|
def contents(self):
|
|
return [filename for filename in self.zf.namelist() if not filename.endswith('/')]
|
|
|
|
def get_content(self, filename):
|
|
return self.zf.read(filename)
|
|
|
|
def write(self, filename=None):
|
|
# fix for Linux zip files read in Windows
|
|
for file in self.zf.filelist:
|
|
file.create_system = 0
|
|
|
|
self.descriptor.seek(0)
|
|
|
|
if filename:
|
|
descriptor = open(filename, 'w')
|
|
descriptor.write(self.descriptor.read())
|
|
else:
|
|
return self.descriptor
|
|
|
|
def as_file(self, filename):
|
|
return SimpleUploadedFile(name=filename, content=self.write().read())
|
|
|
|
def children(self):
|
|
try:
|
|
# Try for a ZIP file
|
|
zfobj = zipfile.ZipFile(self.file_object)
|
|
filenames = [filename for filename in zfobj.namelist() if not filename.endswith('/')]
|
|
return (SimpleUploadedFile(name=filename, content=zfobj.read(filename)) for filename in filenames)
|
|
except zipfile.BadZipfile:
|
|
raise NotACompressedFile
|
|
|
|
def close(self):
|
|
self.zf.close()
|