| #!/usr/bin/python2.4 |
| # |
| # Copyright (C) 2008 Google Inc. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| """Module to compress directories in to series of zip files. |
| |
| This module will take a directory and compress all its contents, including |
| child directories into a series of zip files named N.zip where 'N' ranges from |
| 0 to infinity. The zip files will all be below a certain specified maximum |
| threshold. |
| |
| The directory is compressed with a depth first traversal, each directory's |
| file contents being compressed as it is visisted, before the compression of any |
| child directory's contents. In this way the files within an archive are ordered |
| and the archives themselves are ordered. |
| |
| The class also constructs a 'main.py' file intended for use with Google App |
| Engine with a custom App Engine program not currently distributed with this |
| code base. The custom App Engine runtime can leverage the index files written |
| out by this class to more quickly locate which zip file to serve a given URL |
| from. |
| """ |
| |
| __author__ = 'jmatt@google.com (Justin Mattson)' |
| |
| import optparse |
| import os |
| import stat |
| import sys |
| import zipfile |
| import divide_and_compress_constants |
| |
| |
| def CreateOptionsParser(): |
| """Creates the parser for command line arguments. |
| |
| Returns: |
| A configured optparse.OptionParser object. |
| """ |
| rtn = optparse.OptionParser() |
| rtn.add_option('-s', '--sourcefiles', dest='sourcefiles', default=None, |
| help='The directory containing the files to compress') |
| rtn.add_option('-d', '--destination', dest='destination', default=None, |
| help=('Where to put the archive files, this should not be' |
| ' a child of where the source files exist.')) |
| rtn.add_option('-f', '--filesize', dest='filesize', default='1M', |
| help=('Maximum size of archive files. A number followed by ' |
| 'a magnitude indicator either "B", "K", "M", or "G". ' |
| 'Examples:\n 1000000B == one million BYTES\n' |
| ' 1.2M == one point two MEGABYTES\n' |
| ' 1M == 1048576 BYTES')) |
| rtn.add_option('-n', '--nocompress', action='store_false', dest='compress', |
| default=True, |
| help=('Whether the archive files should be compressed, or ' |
| 'just a concatenation of the source files')) |
| return rtn |
| |
| |
| def VerifyArguments(options, parser): |
| """Runs simple checks on correctness of commandline arguments. |
| |
| Args: |
| options: The command line options passed. |
| parser: The parser object used to parse the command string. |
| """ |
| try: |
| if options.sourcefiles is None or options.destination is None: |
| parser.print_help() |
| sys.exit(-1) |
| except AttributeError: |
| parser.print_help() |
| sys.exit(-1) |
| |
| |
| def ParseSize(size_str): |
| """Parse the file size argument from a string to a number of bytes. |
| |
| Args: |
| size_str: The string representation of the file size. |
| |
| Returns: |
| The file size in bytes. |
| |
| Raises: |
| ValueError: Raises an error if the numeric or qualifier portions of the |
| file size argument is invalid. |
| """ |
| if len(size_str) < 2: |
| raise ValueError(('filesize argument not understood, please include' |
| ' a numeric value and magnitude indicator')) |
| magnitude = size_str[-1] |
| if not magnitude in ('B', 'K', 'M', 'G'): |
| raise ValueError(('filesize magnitude indicator not valid, must be "B",' |
| '"K","M", or "G"')) |
| numeral = float(size_str[:-1]) |
| if magnitude == 'K': |
| numeral *= 1024 |
| elif magnitude == 'M': |
| numeral *= 1048576 |
| elif magnitude == 'G': |
| numeral *= 1073741824 |
| return int(numeral) |
| |
| |
| class DirectoryZipper(object): |
| """Class to compress a directory and all its sub-directories.""" |
| |
| def __init__(self, output_path, base_dir, archive_size, enable_compression): |
| """DirectoryZipper constructor. |
| |
| Args: |
| output_path: A string, the path to write the archives and index file to. |
| base_dir: A string, the directory to compress. |
| archive_size: An number, the maximum size, in bytes, of a single |
| archive file. |
| enable_compression: A boolean, whether or not compression should be |
| enabled, if disabled, the files will be written into an uncompresed |
| zip. |
| """ |
| self.output_dir = output_path |
| self.current_archive = '0.zip' |
| self.base_path = base_dir |
| self.max_size = archive_size |
| self.compress = enable_compression |
| |
| # Set index_fp to None, because we don't know what it will be yet. |
| self.index_fp = None |
| |
| def StartCompress(self): |
| """Start compress of the directory. |
| |
| This will start the compression process and write the archives to the |
| specified output directory. It will also produce an 'index.txt' file in the |
| output directory that maps from file to archive. |
| """ |
| self.index_fp = open(os.path.join(self.output_dir, 'main.py'), 'w') |
| self.index_fp.write(divide_and_compress_constants.file_preamble) |
| os.path.walk(self.base_path, self.CompressDirectory, 1) |
| self.index_fp.write(divide_and_compress_constants.file_endpiece) |
| self.index_fp.close() |
| |
| def RemoveLastFile(self, archive_path=None): |
| """Removes the last item in the archive. |
| |
| This removes the last item in the archive by reading the items out of the |
| archive, adding them to a new archive, deleting the old archive, and |
| moving the new archive to the location of the old archive. |
| |
| Args: |
| archive_path: Path to the archive to modify. This archive should not be |
| open elsewhere, since it will need to be deleted. |
| |
| Returns: |
| A new ZipFile object that points to the modified archive file. |
| """ |
| if archive_path is None: |
| archive_path = os.path.join(self.output_dir, self.current_archive) |
| |
| # Move the old file and create a new one at its old location. |
| root, ext = os.path.splitext(archive_path) |
| old_archive = ''.join([root, '-old', ext]) |
| os.rename(archive_path, old_archive) |
| old_fp = self.OpenZipFileAtPath(old_archive, mode='r') |
| |
| # By default, store uncompressed. |
| compress_bit = zipfile.ZIP_STORED |
| if self.compress: |
| compress_bit = zipfile.ZIP_DEFLATED |
| new_fp = self.OpenZipFileAtPath(archive_path, |
| mode='w', |
| compress=compress_bit) |
| |
| # Read the old archive in a new archive, except the last one. |
| for zip_member in old_fp.infolist()[:-1]: |
| new_fp.writestr(zip_member, old_fp.read(zip_member.filename)) |
| |
| # Close files and delete the old one. |
| old_fp.close() |
| new_fp.close() |
| os.unlink(old_archive) |
| |
| def OpenZipFileAtPath(self, path, mode=None, compress=zipfile.ZIP_DEFLATED): |
| """This method is mainly for testing purposes, eg dependency injection.""" |
| if mode is None: |
| if os.path.exists(path): |
| mode = 'a' |
| else: |
| mode = 'w' |
| |
| if mode == 'r': |
| return zipfile.ZipFile(path, mode) |
| else: |
| return zipfile.ZipFile(path, mode, compress) |
| |
| def CompressDirectory(self, unused_id, dir_path, dir_contents): |
| """Method to compress the given directory. |
| |
| This method compresses the directory 'dir_path'. It will add to an existing |
| zip file that still has space and create new ones as necessary to keep zip |
| file sizes under the maximum specified size. This also writes out the |
| mapping of files to archives to the self.index_fp file descriptor |
| |
| Args: |
| unused_id: A numeric identifier passed by the os.path.walk method, this |
| is not used by this method. |
| dir_path: A string, the path to the directory to compress. |
| dir_contents: A list of directory contents to be compressed. |
| """ |
| # Construct the queue of files to be added that this method will use |
| # it seems that dir_contents is given in reverse alphabetical order, |
| # so put them in alphabetical order by inserting to front of the list. |
| dir_contents.sort() |
| zip_queue = [] |
| for filename in dir_contents: |
| zip_queue.append(os.path.join(dir_path, filename)) |
| compress_bit = zipfile.ZIP_DEFLATED |
| if not self.compress: |
| compress_bit = zipfile.ZIP_STORED |
| |
| # Zip all files in this directory, adding to existing archives and creating |
| # as necessary. |
| while zip_queue: |
| target_file = zip_queue[0] |
| if os.path.isfile(target_file): |
| self.AddFileToArchive(target_file, compress_bit) |
| |
| # See if adding the new file made our archive too large. |
| if not self.ArchiveIsValid(): |
| |
| # IF fixing fails, the last added file was to large, skip it |
| # ELSE the current archive filled normally, make a new one and try |
| # adding the file again. |
| if not self.FixArchive('SIZE'): |
| zip_queue.pop(0) |
| else: |
| self.current_archive = '%i.zip' % ( |
| int(self.current_archive[ |
| 0:self.current_archive.rfind('.zip')]) + 1) |
| else: |
| |
| # Write an index record if necessary. |
| self.WriteIndexRecord() |
| zip_queue.pop(0) |
| else: |
| zip_queue.pop(0) |
| |
| def WriteIndexRecord(self): |
| """Write an index record to the index file. |
| |
| Only write an index record if this is the first file to go into archive |
| |
| Returns: |
| True if an archive record is written, False if it isn't. |
| """ |
| archive = self.OpenZipFileAtPath( |
| os.path.join(self.output_dir, self.current_archive), 'r') |
| archive_index = archive.infolist() |
| if len(archive_index) == 1: |
| self.index_fp.write( |
| '[\'%s\', \'%s\'],\n' % (self.current_archive, |
| archive_index[0].filename)) |
| archive.close() |
| return True |
| else: |
| archive.close() |
| return False |
| |
| def FixArchive(self, problem): |
| """Make the archive compliant. |
| |
| Args: |
| problem: An enum, the reason the archive is invalid. |
| |
| Returns: |
| Whether the file(s) removed to fix the archive could conceivably be |
| in an archive, but for some reason can't be added to this one. |
| """ |
| archive_path = os.path.join(self.output_dir, self.current_archive) |
| return_value = None |
| |
| if problem == 'SIZE': |
| archive_obj = self.OpenZipFileAtPath(archive_path, mode='r') |
| num_archive_files = len(archive_obj.infolist()) |
| |
| # IF there is a single file, that means its too large to compress, |
| # delete the created archive |
| # ELSE do normal finalization. |
| if num_archive_files == 1: |
| print ('WARNING: %s%s is too large to store.' % ( |
| self.base_path, archive_obj.infolist()[0].filename)) |
| archive_obj.close() |
| os.unlink(archive_path) |
| return_value = False |
| else: |
| archive_obj.close() |
| self.RemoveLastFile( |
| os.path.join(self.output_dir, self.current_archive)) |
| print 'Final archive size for %s is %i' % ( |
| self.current_archive, os.path.getsize(archive_path)) |
| return_value = True |
| return return_value |
| |
| def AddFileToArchive(self, filepath, compress_bit): |
| """Add the file at filepath to the current archive. |
| |
| Args: |
| filepath: A string, the path of the file to add. |
| compress_bit: A boolean, whether or not this file should be compressed |
| when added. |
| |
| Returns: |
| True if the file could be added (typically because this is a file) or |
| False if it couldn't be added (typically because its a directory). |
| """ |
| curr_archive_path = os.path.join(self.output_dir, self.current_archive) |
| if os.path.isfile(filepath) and not os.path.islink(filepath): |
| if os.path.getsize(filepath) > 1048576: |
| print 'Warning: %s is potentially too large to serve on GAE' % filepath |
| archive = self.OpenZipFileAtPath(curr_archive_path, |
| compress=compress_bit) |
| # Add the file to the archive. |
| archive.write(filepath, filepath[len(self.base_path):]) |
| archive.close() |
| return True |
| else: |
| return False |
| |
| def ArchiveIsValid(self): |
| """Check whether the archive is valid. |
| |
| Currently this only checks whether the archive is under the required size. |
| The thought is that eventually this will do additional validation |
| |
| Returns: |
| True if the archive is valid, False if its not. |
| """ |
| archive_path = os.path.join(self.output_dir, self.current_archive) |
| return os.path.getsize(archive_path) <= self.max_size |
| |
| |
| def main(argv): |
| parser = CreateOptionsParser() |
| (options, unused_args) = parser.parse_args(args=argv[1:]) |
| VerifyArguments(options, parser) |
| zipper = DirectoryZipper(options.destination, |
| options.sourcefiles, |
| ParseSize(options.filesize), |
| options.compress) |
| zipper.StartCompress() |
| |
| |
| if __name__ == '__main__': |
| main(sys.argv) |
