1# Copyright 2020 The Pigweed Authors 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); you may not 4# use this file except in compliance with the License. You may obtain a copy of 5# the License at 6# 7# https://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT 11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 12# License for the specific language governing permissions and limitations under 13# the License. 14"""Takes a set of input files and zips them up.""" 15 16import argparse 17import pathlib 18import sys 19import zipfile 20 21from collections.abc import Iterable 22 23DEFAULT_DELIMITER = '>' 24 25 26class ZipError(Exception): 27 """Raised when a pw_zip archive can't be built as specified.""" 28 29 30def _parse_args(): 31 parser = argparse.ArgumentParser(description=__doc__) 32 parser.add_argument( 33 '--delimiter', 34 nargs='?', 35 default=DEFAULT_DELIMITER, 36 help='Symbol that separates the path and the zip path destination.', 37 ) 38 parser.add_argument( 39 '--input_list', 40 nargs='+', 41 help='Paths to files and dirs to zip and their desired zip location.', 42 ) 43 parser.add_argument('--out_filename', help='Zip file destination.') 44 45 return parser.parse_args() 46 47 48def zip_up( 49 input_list: Iterable, out_filename: str, delimiter=DEFAULT_DELIMITER 50): 51 """Zips up all input files/dirs. 52 53 Args: 54 input_list: List of strings consisting of file or directory, 55 the delimiter, and a path to the desired .zip destination. 56 out_filename: Path and name of the .zip file. 57 delimiter: string that separates the input source and the zip 58 destination. Defaults to '>'. Examples: 59 '/foo.txt > /' # /foo.txt zipped as /foo.txt 60 '/foo.txt > /bar.txt' # /foo.txt zipped as /bar.txt 61 'foo.txt > /' # foo.txt from invokers dir zipped as /foo.txt 62 '/bar/ > /' # Whole bar dir zipped into / 63 """ 64 with zipfile.ZipFile(out_filename, 'w', zipfile.ZIP_DEFLATED) as zip_file: 65 for _input in input_list: 66 try: 67 source, destination = _input.split(delimiter) 68 source = source.strip() 69 destination = destination.strip() 70 except ValueError as value_error: 71 msg = ( 72 f'Input in the form of "[filename or dir] {delimiter} ' 73 f'/zip_destination/" expected. Instead got:\n {_input}' 74 ) 75 raise ZipError(msg) from value_error 76 if not source: 77 raise ZipError( 78 f'Bad input:\n {_input}\nInput source ' 79 f'cannot be empty. Please specify the input in the form ' 80 f'of "[filename or dir] {delimiter} /zip_destination/".' 81 ) 82 if not destination.startswith('/'): 83 raise ZipError( 84 f'Bad input:\n {_input}\nZip desination ' 85 f'"{destination}" must start with "/" to indicate the ' 86 f'zip file\'s root directory.' 87 ) 88 source_path = pathlib.Path(source) 89 destination_path = pathlib.PurePath(destination) 90 91 # Case: the input source path points to a file. 92 if source_path.is_file(): 93 # Case: "foo.txt > /mydir/"; destination is dir. Put foo.txt 94 # into mydir as /mydir/foo.txt 95 if destination.endswith('/'): 96 zip_file.write( 97 source_path, destination_path / source_path.name 98 ) 99 # Case: "foo.txt > /bar.txt"; destination is a file--rename the 100 # source file: put foo.txt into the zip as /bar.txt 101 else: 102 zip_file.write(source_path, destination_path) 103 continue 104 # Case: the input source path points to a directory. 105 if source_path.is_dir(): 106 zip_up_dir( 107 source, source_path, destination, destination_path, zip_file 108 ) 109 continue 110 raise ZipError(f'Unknown source path\n {source_path}') 111 112 113def zip_up_dir( 114 source: str, 115 source_path: pathlib.Path, 116 destination: str, 117 destination_path: pathlib.PurePath, 118 zip_file: zipfile.ZipFile, 119): 120 if not source.endswith('/'): 121 raise ZipError( 122 f'Source path:\n {source}\nis a directory, but is ' 123 f'missing a trailing "/". The / requirement helps prevent bugs. ' 124 f'To fix, add a trailing /:\n {source}/' 125 ) 126 if not destination.endswith('/'): 127 raise ZipError( 128 f'Destination path:\n {destination}\nis a directory, ' 129 f'but is missing a trailing "/". The / requirement helps prevent ' 130 f'bugs. To fix, add a trailing /:\n {destination}/' 131 ) 132 133 # Walk the directory and add zip all of the files with the 134 # same structure as the source. 135 for file_path in source_path.glob('**/*'): 136 if file_path.is_file(): 137 rel_path = file_path.relative_to(source_path) 138 zip_file.write(file_path, destination_path / rel_path) 139 140 141def main(): 142 zip_up(**vars(_parse_args())) 143 144 145if __name__ == '__main__': 146 try: 147 main() 148 except ZipError as err: 149 print('ERROR:', str(err), file=sys.stderr) 150 sys.exit(1) 151