third_party/pigweed/repo/pw_build/py/pw_build/zip.py

   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
  16 import argparse
  17 import pathlib
  18 import sys
  19 import zipfile
  20
  21 from collections.abc import Iterable
  22
  23 DEFAULT_DELIMITER = '>'
  24
  25
  26 class ZipError(Exception):
  27     """Raised when a pw_zip archive can't be built as specified."""
  28
  29
  30 def _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     parser.add_argument(
  38         '--input_list',
  39         nargs='+',
  40         help='Paths to files and dirs to zip and their desired zip location.')
  41     parser.add_argument('--out_filename', help='Zip file destination.')
  42
  43     return parser.parse_args()
  44
  45
  46 def zip_up(input_list: Iterable,
  47            out_filename: str,
  48            delimiter=DEFAULT_DELIMITER):
  49     """Zips up all input files/dirs.
  50
  51     Args:
  52         input_list: List of strings consisting of file or directory,
  53             the delimiter, and a path to the desired .zip destination.
  54         out_filename: Path and name of the .zip file.
  55         delimiter: string that separates the input source and the zip
  56             destination. Defaults to '>'. Examples:
  57             '/foo.txt > /'         # /foo.txt zipped as /foo.txt
  58             '/foo.txt > /bar.txt'  # /foo.txt zipped as /bar.txt
  59             'foo.txt > /'  # foo.txt from invokers dir zipped as /foo.txt
  60             '/bar/ > /'            # Whole bar dir zipped into /
  61     """
  62     with zipfile.ZipFile(out_filename, 'w', zipfile.ZIP_DEFLATED) as zip_file:
  63         for _input in input_list:
  64             try:
  65                 source, destination = _input.split(delimiter)
  66                 source = source.strip()
  67                 destination = destination.strip()
  68             except ValueError as value_error:
  69                 msg = (
  70                     f'Input in the form of "[filename or dir] {delimiter} '
  71                     f'/zip_destination/" expected. Instead got:\n  {_input}')
  72                 raise ZipError(msg) from value_error
  73             if not source:
  74                 raise ZipError(
  75                     f'Bad input:\n  {_input}\nInput source '
  76                     f'cannot be empty. Please specify the input in the form '
  77                     f'of "[filename or dir] {delimiter} /zip_destination/".')
  78             if not destination.startswith('/'):
  79                 raise ZipError(
  80                     f'Bad input:\n  {_input}\nZip desination '
  81                     f'"{destination}" must start with "/" to indicate the '
  82                     f'zip file\'s root directory.')
  83             source_path = pathlib.Path(source)
  84             destination_path = pathlib.PurePath(destination)
  85
  86             # Case: the input source path points to a file.
  87             if source_path.is_file():
  88                 # Case: "foo.txt > /mydir/"; destination is dir. Put foo.txt
  89                 # into mydir as /mydir/foo.txt
  90                 if destination.endswith('/'):
  91                     zip_file.write(source_path,
  92                                    destination_path / source_path.name)
  93                 # Case: "foo.txt > /bar.txt"; destination is a file--rename the
  94                 # source file: put foo.txt into the zip as /bar.txt
  95                 else:
  96                     zip_file.write(source_path, destination_path)
  97                 continue
  98             # Case: the input source path points to a directory.
  99             if source_path.is_dir():
 100                 zip_up_dir(source, source_path, destination, destination_path,
 101                            zip_file)
 102                 continue
 103             raise ZipError(f'Unknown source path\n  {source_path}')
 104
 105
 106 def zip_up_dir(source: str, source_path: pathlib.Path, destination: str,
 107                destination_path: pathlib.PurePath, zip_file: zipfile.ZipFile):
 108     if not source.endswith('/'):
 109         raise ZipError(
 110             f'Source path:\n  {source}\nis a directory, but is '
 111             f'missing a trailing "/". The / requirement helps prevent bugs. '
 112             f'To fix, add a trailing /:\n  {source}/')
 113     if not destination.endswith('/'):
 114         raise ZipError(
 115             f'Destination path:\n  {destination}\nis a directory, '
 116             f'but is missing a trailing "/". The / requirement helps prevent '
 117             f'bugs. To fix, add a trailing /:\n  {destination}/')
 118
 119     # Walk the directory and add zip all of the files with the
 120     # same structure as the source.
 121     for file_path in source_path.glob('**/*'):
 122         if file_path.is_file():
 123             rel_path = file_path.relative_to(source_path)
 124             zip_file.write(file_path, destination_path / rel_path)
 125
 126
 127 def main():
 128     zip_up(**vars(_parse_args()))
 129
 130
 131 if __name__ == '__main__':
 132     try:
 133         main()
 134     except ZipError as err:
 135         print('ERROR:', str(err), file=sys.stderr)
 136         sys.exit(1)