1#!/usr/bin/python2.4 2# 3# Copyright (C) 2008 Google Inc. 4# 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16# 17 18"""Module to compress directories in to series of zip files. 19 20This module will take a directory and compress all its contents, including 21child directories into a series of zip files named N.zip where 'N' ranges from 220 to infinity. The zip files will all be below a certain specified maximum 23threshold. 24 25The directory is compressed with a depth first traversal, each directory's 26file contents being compressed as it is visisted, before the compression of any 27child directory's contents. In this way the files within an archive are ordered 28and the archives themselves are ordered. 29 30The class also constructs a 'main.py' file intended for use with Google App 31Engine with a custom App Engine program not currently distributed with this 32code base. The custom App Engine runtime can leverage the index files written 33out by this class to more quickly locate which zip file to serve a given URL 34from. 35""" 36 37__author__ = 'jmatt@google.com (Justin Mattson)' 38 39import optparse 40import os 41import stat 42import sys 43import zipfile 44import divide_and_compress_constants 45 46 47def CreateOptionsParser(): 48 """Creates the parser for command line arguments. 49 50 Returns: 51 A configured optparse.OptionParser object. 52 """ 53 rtn = optparse.OptionParser() 54 rtn.add_option('-s', '--sourcefiles', dest='sourcefiles', default=None, 55 help='The directory containing the files to compress') 56 rtn.add_option('-d', '--destination', dest='destination', default=None, 57 help=('Where to put the archive files, this should not be' 58 ' a child of where the source files exist.')) 59 rtn.add_option('-f', '--filesize', dest='filesize', default='1M', 60 help=('Maximum size of archive files. A number followed by ' 61 'a magnitude indicator either "B", "K", "M", or "G". ' 62 'Examples:\n 1000000B == one million BYTES\n' 63 ' 1.2M == one point two MEGABYTES\n' 64 ' 1M == 1048576 BYTES')) 65 rtn.add_option('-n', '--nocompress', action='store_false', dest='compress', 66 default=True, 67 help=('Whether the archive files should be compressed, or ' 68 'just a concatenation of the source files')) 69 return rtn 70 71 72def VerifyArguments(options, parser): 73 """Runs simple checks on correctness of commandline arguments. 74 75 Args: 76 options: The command line options passed. 77 parser: The parser object used to parse the command string. 78 """ 79 try: 80 if options.sourcefiles is None or options.destination is None: 81 parser.print_help() 82 sys.exit(-1) 83 except AttributeError: 84 parser.print_help() 85 sys.exit(-1) 86 87 88def ParseSize(size_str): 89 """Parse the file size argument from a string to a number of bytes. 90 91 Args: 92 size_str: The string representation of the file size. 93 94 Returns: 95 The file size in bytes. 96 97 Raises: 98 ValueError: Raises an error if the numeric or qualifier portions of the 99 file size argument is invalid. 100 """ 101 if len(size_str) < 2: 102 raise ValueError(('filesize argument not understood, please include' 103 ' a numeric value and magnitude indicator')) 104 magnitude = size_str[-1] 105 if not magnitude in ('B', 'K', 'M', 'G'): 106 raise ValueError(('filesize magnitude indicator not valid, must be "B",' 107 '"K","M", or "G"')) 108 numeral = float(size_str[:-1]) 109 if magnitude == 'K': 110 numeral *= 1024 111 elif magnitude == 'M': 112 numeral *= 1048576 113 elif magnitude == 'G': 114 numeral *= 1073741824 115 return int(numeral) 116 117 118class DirectoryZipper(object): 119 """Class to compress a directory and all its sub-directories.""" 120 121 def __init__(self, output_path, base_dir, archive_size, enable_compression): 122 """DirectoryZipper constructor. 123 124 Args: 125 output_path: A string, the path to write the archives and index file to. 126 base_dir: A string, the directory to compress. 127 archive_size: An number, the maximum size, in bytes, of a single 128 archive file. 129 enable_compression: A boolean, whether or not compression should be 130 enabled, if disabled, the files will be written into an uncompresed 131 zip. 132 """ 133 self.output_dir = output_path 134 self.current_archive = '0.zip' 135 self.base_path = base_dir 136 self.max_size = archive_size 137 self.compress = enable_compression 138 139 # Set index_fp to None, because we don't know what it will be yet. 140 self.index_fp = None 141 142 def StartCompress(self): 143 """Start compress of the directory. 144 145 This will start the compression process and write the archives to the 146 specified output directory. It will also produce an 'index.txt' file in the 147 output directory that maps from file to archive. 148 """ 149 self.index_fp = open(os.path.join(self.output_dir, 'main.py'), 'w') 150 self.index_fp.write(divide_and_compress_constants.file_preamble) 151 os.path.walk(self.base_path, self.CompressDirectory, 1) 152 self.index_fp.write(divide_and_compress_constants.file_endpiece) 153 self.index_fp.close() 154 155 def RemoveLastFile(self, archive_path=None): 156 """Removes the last item in the archive. 157 158 This removes the last item in the archive by reading the items out of the 159 archive, adding them to a new archive, deleting the old archive, and 160 moving the new archive to the location of the old archive. 161 162 Args: 163 archive_path: Path to the archive to modify. This archive should not be 164 open elsewhere, since it will need to be deleted. 165 166 Returns: 167 A new ZipFile object that points to the modified archive file. 168 """ 169 if archive_path is None: 170 archive_path = os.path.join(self.output_dir, self.current_archive) 171 172 # Move the old file and create a new one at its old location. 173 root, ext = os.path.splitext(archive_path) 174 old_archive = ''.join([root, '-old', ext]) 175 os.rename(archive_path, old_archive) 176 old_fp = self.OpenZipFileAtPath(old_archive, mode='r') 177 178 # By default, store uncompressed. 179 compress_bit = zipfile.ZIP_STORED 180 if self.compress: 181 compress_bit = zipfile.ZIP_DEFLATED 182 new_fp = self.OpenZipFileAtPath(archive_path, 183 mode='w', 184 compress=compress_bit) 185 186 # Read the old archive in a new archive, except the last one. 187 for zip_member in old_fp.infolist()[:-1]: 188 new_fp.writestr(zip_member, old_fp.read(zip_member.filename)) 189 190 # Close files and delete the old one. 191 old_fp.close() 192 new_fp.close() 193 os.unlink(old_archive) 194 195 def OpenZipFileAtPath(self, path, mode=None, compress=zipfile.ZIP_DEFLATED): 196 """This method is mainly for testing purposes, eg dependency injection.""" 197 if mode is None: 198 if os.path.exists(path): 199 mode = 'a' 200 else: 201 mode = 'w' 202 203 if mode == 'r': 204 return zipfile.ZipFile(path, mode) 205 else: 206 return zipfile.ZipFile(path, mode, compress) 207 208 def CompressDirectory(self, unused_id, dir_path, dir_contents): 209 """Method to compress the given directory. 210 211 This method compresses the directory 'dir_path'. It will add to an existing 212 zip file that still has space and create new ones as necessary to keep zip 213 file sizes under the maximum specified size. This also writes out the 214 mapping of files to archives to the self.index_fp file descriptor 215 216 Args: 217 unused_id: A numeric identifier passed by the os.path.walk method, this 218 is not used by this method. 219 dir_path: A string, the path to the directory to compress. 220 dir_contents: A list of directory contents to be compressed. 221 """ 222 # Construct the queue of files to be added that this method will use 223 # it seems that dir_contents is given in reverse alphabetical order, 224 # so put them in alphabetical order by inserting to front of the list. 225 dir_contents.sort() 226 zip_queue = [] 227 for filename in dir_contents: 228 zip_queue.append(os.path.join(dir_path, filename)) 229 compress_bit = zipfile.ZIP_DEFLATED 230 if not self.compress: 231 compress_bit = zipfile.ZIP_STORED 232 233 # Zip all files in this directory, adding to existing archives and creating 234 # as necessary. 235 while zip_queue: 236 target_file = zip_queue[0] 237 if os.path.isfile(target_file): 238 self.AddFileToArchive(target_file, compress_bit) 239 240 # See if adding the new file made our archive too large. 241 if not self.ArchiveIsValid(): 242 243 # IF fixing fails, the last added file was to large, skip it 244 # ELSE the current archive filled normally, make a new one and try 245 # adding the file again. 246 if not self.FixArchive('SIZE'): 247 zip_queue.pop(0) 248 else: 249 self.current_archive = '%i.zip' % ( 250 int(self.current_archive[ 251 0:self.current_archive.rfind('.zip')]) + 1) 252 else: 253 254 # Write an index record if necessary. 255 self.WriteIndexRecord() 256 zip_queue.pop(0) 257 else: 258 zip_queue.pop(0) 259 260 def WriteIndexRecord(self): 261 """Write an index record to the index file. 262 263 Only write an index record if this is the first file to go into archive 264 265 Returns: 266 True if an archive record is written, False if it isn't. 267 """ 268 archive = self.OpenZipFileAtPath( 269 os.path.join(self.output_dir, self.current_archive), 'r') 270 archive_index = archive.infolist() 271 if len(archive_index) == 1: 272 self.index_fp.write( 273 '[\'%s\', \'%s\'],\n' % (self.current_archive, 274 archive_index[0].filename)) 275 archive.close() 276 return True 277 else: 278 archive.close() 279 return False 280 281 def FixArchive(self, problem): 282 """Make the archive compliant. 283 284 Args: 285 problem: An enum, the reason the archive is invalid. 286 287 Returns: 288 Whether the file(s) removed to fix the archive could conceivably be 289 in an archive, but for some reason can't be added to this one. 290 """ 291 archive_path = os.path.join(self.output_dir, self.current_archive) 292 return_value = None 293 294 if problem == 'SIZE': 295 archive_obj = self.OpenZipFileAtPath(archive_path, mode='r') 296 num_archive_files = len(archive_obj.infolist()) 297 298 # IF there is a single file, that means its too large to compress, 299 # delete the created archive 300 # ELSE do normal finalization. 301 if num_archive_files == 1: 302 print ('WARNING: %s%s is too large to store.' % ( 303 self.base_path, archive_obj.infolist()[0].filename)) 304 archive_obj.close() 305 os.unlink(archive_path) 306 return_value = False 307 else: 308 archive_obj.close() 309 self.RemoveLastFile( 310 os.path.join(self.output_dir, self.current_archive)) 311 print 'Final archive size for %s is %i' % ( 312 self.current_archive, os.path.getsize(archive_path)) 313 return_value = True 314 return return_value 315 316 def AddFileToArchive(self, filepath, compress_bit): 317 """Add the file at filepath to the current archive. 318 319 Args: 320 filepath: A string, the path of the file to add. 321 compress_bit: A boolean, whether or not this file should be compressed 322 when added. 323 324 Returns: 325 True if the file could be added (typically because this is a file) or 326 False if it couldn't be added (typically because its a directory). 327 """ 328 curr_archive_path = os.path.join(self.output_dir, self.current_archive) 329 if os.path.isfile(filepath) and not os.path.islink(filepath): 330 if os.path.getsize(filepath) > 1048576: 331 print 'Warning: %s is potentially too large to serve on GAE' % filepath 332 archive = self.OpenZipFileAtPath(curr_archive_path, 333 compress=compress_bit) 334 # Add the file to the archive. 335 archive.write(filepath, filepath[len(self.base_path):]) 336 archive.close() 337 return True 338 else: 339 return False 340 341 def ArchiveIsValid(self): 342 """Check whether the archive is valid. 343 344 Currently this only checks whether the archive is under the required size. 345 The thought is that eventually this will do additional validation 346 347 Returns: 348 True if the archive is valid, False if its not. 349 """ 350 archive_path = os.path.join(self.output_dir, self.current_archive) 351 return os.path.getsize(archive_path) <= self.max_size 352 353 354def main(argv): 355 parser = CreateOptionsParser() 356 (options, unused_args) = parser.parse_args(args=argv[1:]) 357 VerifyArguments(options, parser) 358 zipper = DirectoryZipper(options.destination, 359 options.sourcefiles, 360 ParseSize(options.filesize), 361 options.compress) 362 zipper.StartCompress() 363 364 365if __name__ == '__main__': 366 main(sys.argv) 367