diff --git a/pip.py b/pip.py index 94ef7e887..30d2edcf0 100644 --- a/pip.py +++ b/pip.py @@ -452,6 +452,292 @@ class FreezeCommand(Command): FreezeCommand() +class ZipCommand(Command): + name = 'zip' + usage = '%prog [OPTIONS] PACKAGE_NAMES...' + summary = 'zip (or unzip) individual packages' + + def __init__(self): + super(ZipCommand, self).__init__() + if self.name == 'zip': + self.parser.add_option( + '--unzip', + action='store_true', + dest='unzip', + help='Unzip (rather than zip) a package') + else: + self.parser.add_option( + '--zip', + action='store_false', + dest='unzip', + default=True, + help='Zip (rather than unzip) a package') + self.parser.add_option( + '--no-pyc', + action='store_true', + dest='no_pyc', + help='Do not include .pyc files in zip files (useful on Google App Engine)') + self.parser.add_option( + '-l', '--list', + action='store_true', + dest='list', + help='List the packages available, and their zip status') + self.parser.add_option( + '--sort-files', + action='store_true', + dest='sort_files', + help='With --list, sort packages according to how many files they contain') + + def run(self, options, args): + if options.list: + return self.list(options, args) + if not args: + raise InstallationError( + 'You must give at least one package to zip or unzip') + packages = [] + for arg in args: + module_name, filename = self.find_package(arg) + if options.unzip and os.path.isdir(filename): + raise InstallationError( + 'The module %s (in %s) is not a zip file; cannot be unzipped' + % (module_name, filename)) + elif not options.unzip and not os.path.isdir(filename): + raise InstallationError( + 'The module %s (in %s) is not a directory; cannot be zipped' + % (module_name, filename)) + packages.append((module_name, filename)) + for module_name, filename in packages: + if options.unzip: + return self.unzip_package(module_name, filename) + else: + return self.zip_package(module_name, filename, options.no_pyc) + + def unzip_package(self, module_name, filename): + zip_filename = os.path.dirname(filename) + if not os.path.isfile(zip_filename) and zipfile.is_zipfile(zip_filename): + raise InstallationError( + 'Module %s (in %s) isn\'t located in a zip file in %s' + % (module_name, filename, zip_filename)) + package_path = os.path.dirname(zip_filename) + if not package_path in sys.path: + logger.warn( + 'Unpacking %s into %s, but %s is not on sys.path' + % (display_path(zip_filename), display_path(package_path), + display_path(package_path))) + logger.notify('Unzipping %s (in %s)' % (module_name, display_path(zip_filename))) + logger.indent += 2 + try: + ## FIXME: this should be undoable: + zip = zipfile.ZipFile(zip_filename) + to_save = [] + for name in zip.namelist(): + if name.startswith('%s/' % module_name): + content = zip.read(name) + dest = os.path.join(package_path, name) + if not os.path.exists(os.path.dirname(dest)): + os.makedirs(os.path.dirname(dest)) + if not content and dest.endswith('/'): + if not os.path.exists(dest): + os.makedirs(dest) + else: + f = open(dest, 'wb') + f.write(content) + f.close() + else: + to_save.append((name, zip.read(name))) + zip.close() + if not to_save: + logger.info('Removing now-empty zip file %s' % display_path(zip_filename)) + os.unlink(zip_filename) + self.remove_filename_from_pth(zip_filename) + else: + logger.info('Removing entries in %s/ from zip file %s' % (module_name, display_path(zip_filename))) + zip = zipfile.ZipFile(zip_filename, 'w') + for name, content in to_save: + zip.writestr(name, content) + zip.close() + finally: + logger.indent -= 2 + + def zip_package(self, module_name, filename, no_pyc): + orig_filename = filename + logger.notify('Zip %s (in %s)' % (module_name, display_path(filename))) + logger.indent += 2 + if filename.endswith('.egg'): + dest_filename = filename + else: + dest_filename = filename + '.zip' + try: + ## FIXME: I think this needs to be undoable: + if filename == dest_filename: + filename = backup_dir(orig_filename) + shutil.move(orig_filename, filename) + try: + logger.info('Creating zip file in %s' % display_path(dest_filename)) + zip = zipfile.ZipFile(dest_filename, 'w') + zip.writestr(module_name + '/', '') + for dirpath, dirnames, filenames in os.walk(filename): + if no_pyc: + filenames = [f for f in filenames + if not f.lower().endswith('.pyc')] + for fns, is_dir in [(dirnames, True), (filenames, False)]: + for fn in fns: + full = os.path.join(dirpath, fn) + dest = os.path.join(module_name, dirpath[len(filename):].lstrip(os.path.sep), fn) + if is_dir: + zip.writestr(dest+'/', '') + else: + zip.write(full, dest) + zip.close() + logger.info('Removing old directory %s' % display_path(filename)) + shutil.rmtree(filename) + except: + ## FIXME: need to do an undo here + raise + ## FIXME: should also be undone; + self.add_filename_to_pth(dest_filename) + finally: + logger.indent -= 2 + + def remove_filename_from_pth(self, filename): + for pth in self.pth_files(): + f = open(pth, 'r') + lines = f.readlines() + f.close() + new_lines = [ + l for l in lines if l.strip() != filename] + if lines != new_lines: + logger.info('Removing reference to %s from .pth file %s' + % (display_path(filename), display_path(pth))) + if not filter(None, new_lines): + logger.info('%s file would be empty: deleting' % display_path(pth)) + os.unlink(pth) + else: + f = open(pth, 'w') + f.writelines(new_lines) + f.close() + return + logger.warn('Cannot find a reference to %s in any .pth file' % display_path(filename)) + + def add_filename_to_pth(self, filename): + path = os.path.dirname(filename) + dest = os.path.join(path, filename + '.pth') + if path not in sys.path: + logger.warn('Adding .pth file %s, but it is not on sys.path' % display_path(dest)) + if os.path.exists(dest): + f = open(dest) + lines = f.readlines() + f.close() + if lines and not lines[-1].endswith('\n'): + lines[-1] += '\n' + lines.append(filename+'\n') + else: + lines = [filename + '\n'] + f = open(dest, 'w') + f.writelines(lines) + f.close() + + def pth_files(self): + for path in sys.path: + if not os.path.exists(path) or not os.path.isdir(path): + continue + for filename in os.listdir(path): + if filename.endswith('.pth'): + yield os.path.join(path, filename) + + def find_package(self, package): + for path in sys.path: + full = os.path.join(path, package) + if os.path.exists(full): + return package, full + if not os.path.isdir(path) and zipfile.is_zipfile(path): + zip = zipfile.ZipFile(path, 'r') + try: + zip.read('%s/__init__.py' % package) + except KeyError: + pass + else: + zip.close() + return package, full + zip.close() + ## FIXME: need special error for package.py case: + raise InstallationError( + 'No package with the name %s found' % package) + + def list(self, options, args): + if args: + raise InstallationError( + 'You cannot give an argument with --list') + for path in sorted(sys.path): + if not os.path.exists(path): + continue + basename = os.path.basename(path.rstrip(os.path.sep)) + if os.path.isfile(path) and zipfile.is_zipfile(path): + if os.path.dirname(path) not in sys.path: + logger.notify('Zipped egg: %s' % display_path(path)) + continue + if basename != 'site-packages': + continue + logger.notify('In %s:' % display_path(path)) + logger.indent += 2 + zipped = [] + unzipped = [] + try: + for filename in sorted(os.listdir(path)): + ext = os.path.splitext(filename)[1].lower() + if ext in ('.pth', '.egg-info', '.egg-link'): + continue + if ext == '.py': + logger.info('Not displaying %s: not a package' % display_path(filename)) + continue + full = os.path.join(path, filename) + if os.path.isdir(full): + unzipped.append((filename, self.count_package(full))) + elif zipfile.is_zipfile(full): + zipped.append(filename) + else: + logger.info('Unknown file: %s' % display_path(filename)) + if zipped: + logger.notify('Zipped packages:') + logger.indent += 2 + try: + for filename in zipped: + logger.notify(filename) + finally: + logger.indent -= 2 + else: + logger.notify('No zipped packages.') + if unzipped: + if options.sort_files: + unzipped.sort(key=lambda x: -x[1]) + logger.notify('Unzipped packages:') + logger.indent += 2 + try: + for filename, count in unzipped: + logger.notify('%s (%i files)' % (filename, count)) + finally: + logger.indent -= 2 + else: + logger.notify('No unzipped packages.') + finally: + logger.indent -= 2 + + def count_package(self, path): + total = 0 + for dirpath, dirnames, filenames in os.walk(path): + filenames = [f for f in filenames + if not f.lower().endswith('.pyc')] + total += len(filenames) + return total + +ZipCommand() + +class UnzipCommand(ZipCommand): + name = 'unzip' + +UnzipCommand() + + def main(initial_args=None): if initial_args is None: initial_args = sys.argv[1:]