| |
| from __future__ import with_statement |
| import logging |
| import os.path |
| from zipfile import ZIP_DEFLATED |
| from path import split as path_split |
| from path import get_ancestors as path_ancestors |
|
|
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def zipfile_nodes(zipfile): |
| seen = set() |
| for path in zipfile.namelist(): |
| for anc_path in path_ancestors(path): |
| if anc_path not in seen: |
| yield anc_path, ZipFileFolder(zipfile, anc_path) |
| seen.add(anc_path) |
| if path not in seen: |
| if path.endswith('/'): |
| yield path, ZipFileFolder(zipfile, path) |
| else: |
| yield path, ZipFileFile(zipfile, path) |
| seen.add(path) |
|
|
|
|
| class ZipFileNode(object): |
| |
| def __init__(self, zipfile, path): |
| self.zipfile = zipfile |
| self.path = path |
|
|
|
|
| class ZipFileFile(ZipFileNode): |
|
|
| def open(self, mode='r', compress_type=ZIP_DEFLATED): |
| if mode.startswith('r'): |
| return ZipFileStream(self.zipfile.open(self.path, mode)) |
| elif mode == 'w': |
| return ZipFileStream(ZipFileWritable(self.zipfile, self.path, |
| compress_type=compress_type)) |
|
|
| def put(self, filesystem_path, compress_type=ZIP_DEFLATED): |
| self.zipfile.write(filesystem_path, self.path, compress_type) |
|
|
|
|
| class ZipFileFolder(ZipFileNode): |
|
|
| def childs(self): |
| prefix = path_split(self.path) |
| prefix_len = len(prefix) |
| for path, node in zipfile_nodes(self.zipfile): |
| path_segments = path_split(path) |
| if len(path_segments) == prefix_len + 1: |
| if path_segments[:prefix_len] == prefix: |
| yield path_segments[-1], node |
|
|
| def __iter__(self): |
| for name, node in self.childs(): |
| yield name |
|
|
| def __getitem__(self, name): |
| for node_name, node in self.childs(): |
| if node_name == name: |
| return node |
| raise KeyError(name) |
|
|
| def file(self, name): |
| path = os.path.join(self.path, name) |
| return ZipFileFile(self.zipfile, path) |
|
|
| def folder(self, name): |
| path = os.path.join(self.path, name) |
| return ZipFileFolder(self.zipfile, path) |
|
|
|
|
| class ZipFileStorage(ZipFileFolder): |
|
|
| def __init__(self, *args, **kwargs): |
| import zipfile |
| self.zipfile = zipfile.ZipFile(*args, **kwargs) |
| self.path = '' |
|
|
| def childs(self): |
| for path, node in zipfile_nodes(self.zipfile): |
| path_segments = path_split(path) |
| if len(path_segments) == 1: |
| yield path_segments[-1], node |
|
|
| def close(self): |
| self.zipfile.close() |
|
|
| def __enter__(self): |
| return self |
|
|
| def __exit__(self, *args, **kwargs): |
| self.close() |
|
|
|
|
| class ZipFileStream(object): |
|
|
| def __init__(self, stream): |
| self.stream = stream |
|
|
| def __getattr__(self, name): |
| return getattr(self.stream, name) |
|
|
| def __enter__(self): |
| return self.stream |
|
|
| def __exit__(self, *args, **kwargs): |
| self.close() |
|
|
|
|
| class ZipFileWritable(object): |
|
|
| def __init__(self, zipfile, path, compress_type=ZIP_DEFLATED): |
| self.zipfile = zipfile |
| self.path = path |
| self.compress_type = compress_type |
|
|
| import tempfile |
| fd, tmp_path = tempfile.mkstemp() |
| self.tmp_f = os.fdopen(fd, 'w') |
| self.tmp_path = tmp_path |
|
|
| def __getattr__(self, name): |
| return getattr(self.tmp_f, name) |
|
|
| def close(self): |
| self.tmp_f.close() |
| self.zipfile.write(self.tmp_path, self.path, self.compress_type) |
| os.unlink(self.tmp_path) |
|
|
