Index: swarm_client/utils/zip_package.py |
=================================================================== |
--- swarm_client/utils/zip_package.py (revision 235167) |
+++ swarm_client/utils/zip_package.py (working copy) |
@@ -1,292 +0,0 @@ |
-# Copyright 2013 The Chromium Authors. All rights reserved. |
-# Use of this source code is governed by a BSD-style license that can be |
-# found in the LICENSE file. |
- |
-"""Utilities to work with importable python zip packages.""" |
- |
-import atexit |
-import collections |
-import cStringIO as StringIO |
-import os |
-import pkgutil |
-import re |
-import sys |
-import tempfile |
-import threading |
-import zipfile |
-import zipimport |
- |
- |
-# Glob patterns for files to exclude from a package by default. |
-EXCLUDE_LIST = ( |
- # Ignore hidden files (including .svn and .git). |
- r'\..*', |
- |
- # Ignore precompiled python files since they depend on python version and we |
- # don't want zip package to be version-depended. |
- r'.*\.pyc$', |
- r'.*\.pyo$', |
-) |
- |
- |
-# Temporary files extracted by extract_resource. Removed in atexit hook. |
-_extracted_files = [] |
-_extracted_files_lock = threading.Lock() |
- |
- |
-class ZipPackageError(RuntimeError): |
- """Failed to create a zip package.""" |
- |
- |
-class ZipPackage(object): |
- """A set of files that can be zipped to file on disk or into memory buffer. |
- |
- Usage: |
- package = ZipPackage(root) |
- package.add_file('some_file.py', '__main__.py') |
- package.add_directory('some_directory') |
- package.add_buffer('generated.py', 'any string here') |
- |
- buf = package.zip_into_buffer() |
- package.zip_into_file('my_zip.zip') |
- """ |
- |
- _FileRef = collections.namedtuple('_FileRef', ['abs_path']) |
- _BufferRef = collections.namedtuple('_BufferRef', ['buffer']) |
- |
- def __init__(self, root): |
- """Initializes new empty ZipPackage. |
- |
- All files added to the package should live under the |root|. It will also |
- be used when calculating relative paths of files in the package. |
- |
- |root| must be an absolute path. |
- """ |
- assert os.path.isabs(root), root |
- self.root = root.rstrip(os.sep) + os.sep |
- self._items = {} |
- |
- @property |
- def files(self): |
- """Files added to the package as a list of relative paths in zip.""" |
- return self._items.keys() |
- |
- def add_file(self, absolute_path, archive_path=None): |
- """Adds a single file to the package. |
- |
- |archive_path| is a relative path in archive for this file, by default it's |
- equal to |absolute_path| taken relative to |root|. In that case |
- |absolute_path| must be in a |root| subtree. |
- |
- If |archive_path| is given, |absolute_path| can point to any file. |
- """ |
- assert os.path.isabs(absolute_path), absolute_path |
- absolute_path = os.path.normpath(absolute_path) |
- # If |archive_path| is not given, ensure that |absolute_path| is under root. |
- if not archive_path and not absolute_path.startswith(self.root): |
- raise ZipPackageError( |
- 'Path %s is not inside root %s' % (absolute_path, self.root)) |
- if not os.path.exists(absolute_path): |
- raise ZipPackageError('No such file: %s' % absolute_path) |
- if not os.path.isfile(absolute_path): |
- raise ZipPackageError('Object %s is not a regular file' % absolute_path) |
- archive_path = archive_path or absolute_path[len(self.root):] |
- self._add_entry(archive_path, ZipPackage._FileRef(absolute_path)) |
- |
- def add_python_file(self, absolute_path, archive_path=None): |
- """Adds a single python file to the package. |
- |
- Recognizes *.pyc files and adds corresponding *.py file instead. |
- """ |
- base, ext = os.path.splitext(absolute_path) |
- if ext in ('.pyc', '.pyo'): |
- absolute_path = base + '.py' |
- elif ext != '.py': |
- raise ZipPackageError('Not a python file: %s' % absolute_path) |
- self.add_file(absolute_path, archive_path) |
- |
- def add_directory(self, absolute_path, archive_path=None, |
- exclude=EXCLUDE_LIST): |
- """Recursively adds all files from given directory to the package. |
- |
- |archive_path| is a relative path in archive for this directory, by default |
- it's equal to |absolute_path| taken relative to |root|. In that case |
- |absolute_path| must be in |root| subtree. |
- |
- If |archive_path| is given, |absolute_path| can point to any directory. |
- |
- |exclude| defines a list of regular expressions for file names to exclude |
- from the package. |
- |
- Only non-empty directories will be actually added to the package. |
- """ |
- assert os.path.isabs(absolute_path), absolute_path |
- absolute_path = os.path.normpath(absolute_path).rstrip(os.sep) + os.sep |
- # If |archive_path| is not given, ensure that |path| is under root. |
- if not archive_path and not absolute_path.startswith(self.root): |
- raise ZipPackageError( |
- 'Path %s is not inside root %s' % (absolute_path, self.root)) |
- if not os.path.exists(absolute_path): |
- raise ZipPackageError('No such directory: %s' % absolute_path) |
- if not os.path.isdir(absolute_path): |
- raise ZipPackageError('Object %s is not a directory' % absolute_path) |
- |
- # Precompile regular expressions. |
- exclude_regexps = [re.compile(r) for r in exclude] |
- # Returns True if |name| should be excluded from the package. |
- should_exclude = lambda name: any(r.match(name) for r in exclude_regexps) |
- |
- archive_path = archive_path or absolute_path[len(self.root):] |
- for cur_dir, dirs, files in os.walk(absolute_path): |
- # Add all non-excluded files. |
- for name in files: |
- if not should_exclude(name): |
- absolute = os.path.join(cur_dir, name) |
- relative = absolute[len(absolute_path):] |
- assert absolute.startswith(absolute_path) |
- self.add_file(absolute, os.path.join(archive_path, relative)) |
- # Remove excluded directories from enumeration. |
- for name in [d for d in dirs if should_exclude(d)]: |
- dirs.remove(name) |
- |
- def add_buffer(self, archive_path, buf): |
- """Adds a contents of the given string |buf| to the package as a file. |
- |
- |archive_path| is a path in archive for this file. |
- """ |
- # Only 'str' is allowed here, no 'unicode' |
- assert isinstance(buf, str) |
- self._add_entry(archive_path, ZipPackage._BufferRef(buf)) |
- |
- def zip_into_buffer(self, compress=True): |
- """Zips added files into in-memory zip file and returns it as str.""" |
- stream = StringIO.StringIO() |
- try: |
- self._zip_into_stream(stream, compress) |
- return stream.getvalue() |
- finally: |
- stream.close() |
- |
- def zip_into_file(self, path, compress=True): |
- """Zips added files into a file on disk.""" |
- with open(path, 'wb') as stream: |
- self._zip_into_stream(stream, compress) |
- |
- def _add_entry(self, archive_path, ref): |
- """Adds new zip package entry.""" |
- # Always use forward slashes in zip. |
- archive_path = archive_path.replace(os.sep, '/') |
- # Ensure there are no suspicious components in the path. |
- assert not any(p in ('', '.', '..') for p in archive_path.split('/')) |
- # Ensure there's no file overwrites. |
- if archive_path in self._items: |
- raise ZipPackageError('Duplicated entry: %s' % archive_path) |
- self._items[archive_path] = ref |
- |
- def _zip_into_stream(self, stream, compress): |
- """Zips files added so far into some output stream. |
- |
- Some measures are taken to guarantee that final zip depends only on the |
- content of added files: |
- * File modification time is not stored. |
- * Entries are sorted by file name in archive. |
- """ |
- compression = zipfile.ZIP_DEFLATED if compress else zipfile.ZIP_STORED |
- zip_file = zipfile.ZipFile(stream, 'w', compression) |
- try: |
- for archive_path in sorted(self._items): |
- ref = self._items[archive_path] |
- info = zipfile.ZipInfo(filename=archive_path) |
- info.compress_type = compression |
- info.create_system = 3 |
- if isinstance(ref, ZipPackage._FileRef): |
- info.external_attr = (os.stat(ref.abs_path)[0] & 0xFFFF) << 16L |
- with open(ref.abs_path, 'rb') as f: |
- buf = f.read() |
- elif isinstance(ref, ZipPackage._BufferRef): |
- buf = ref.buffer |
- else: |
- assert False, 'Unexpected type %s' % ref |
- zip_file.writestr(info, buf) |
- finally: |
- zip_file.close() |
- |
- |
-def get_module_zip_archive(module): |
- """Given a module, returns path to a zip package that contains it or None.""" |
- loader = pkgutil.get_loader(module) |
- if not isinstance(loader, zipimport.zipimporter): |
- return None |
- # 'archive' property is documented only for python 2.7, but it appears to be |
- # there at least since python 2.5.2. |
- return loader.archive |
- |
- |
-def is_zipped_module(module): |
- """True if given module was loaded from a zip package.""" |
- return bool(get_module_zip_archive(module)) |
- |
- |
-def get_main_script_path(): |
- """If running from zip returns path to a zip file, else path to __main__. |
- |
- Basically returns path to a file passed to python for execution |
- as in 'python <main_script>' considering a case of executable zip package. |
- |
- Returns path relative to a current directory of when process was started. |
- """ |
- # If running from interactive console __file__ is not defined. |
- main = sys.modules['__main__'] |
- return get_module_zip_archive(main) or getattr(main, '__file__', None) |
- |
- |
-def extract_resource(package, resource): |
- """Returns real file system path to a |resource| file from a |package|. |
- |
- If it's inside a zip package, will extract it first into temp file created |
- with tempfile.mkstemp. Such file is readable and writable only by the creating |
- user ID. |
- |
- |package| is a python module object that represents a package. |
- |resource| should be a relative filename, using '/'' as the path separator. |
- |
- Raises ValueError if no such resource. |
- """ |
- # For regular non-zip packages just construct an absolute path. |
- if not is_zipped_module(package): |
- # Package's __file__ attribute is always an absolute path. |
- path = os.path.join(os.path.dirname(package.__file__), |
- resource.replace('/', os.sep)) |
- if not os.path.exists(path): |
- raise ValueError('No such resource in %s: %s' % (package, resource)) |
- return path |
- |
- # For zipped packages extract the resource into a temp file. |
- data = pkgutil.get_data(package.__name__, resource) |
- if data is None: |
- raise ValueError('No such resource in zipped %s: %s' % (package, resource)) |
- fd, path = tempfile.mkstemp() |
- with os.fdopen(fd, 'w') as stream: |
- stream.write(data) |
- |
- # Register it for removal when process dies. |
- with _extracted_files_lock: |
- _extracted_files.append(path) |
- # First extracted file -> register atexit hook that cleans them all. |
- if len(_extracted_files) == 1: |
- atexit.register(cleanup_extracted_resources) |
- |
- return path |
- |
- |
-def cleanup_extracted_resources(): |
- """Removes all temporary files created by extract_resource. |
- |
- Executed as atexit hook. |
- """ |
- with _extracted_files_lock: |
- while _extracted_files: |
- try: |
- os.remove(_extracted_files.pop()) |
- except OSError: |
- pass |