Christopher Dunn | 9cc0bb8 | 2015-01-24 15:44:51 -0600 | [diff] [blame] | 1 | from contextlib import closing |
Christopher Dunn | f357688 | 2015-01-24 16:20:25 -0600 | [diff] [blame] | 2 | import os |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 3 | import tarfile |
| 4 | |
| 5 | TARGZ_DEFAULT_COMPRESSION_LEVEL = 9 |
| 6 | |
| 7 | def make_tarball(tarball_path, sources, base_dir, prefix_dir=''): |
| 8 | """Parameters: |
| 9 | tarball_path: output path of the .tar.gz file |
| 10 | sources: list of sources to include in the tarball, relative to the current directory |
| 11 | base_dir: if a source file is in a sub-directory of base_dir, then base_dir is stripped |
| 12 | from path in the tarball. |
| 13 | prefix_dir: all files stored in the tarball be sub-directory of prefix_dir. Set to '' |
| 14 | to make them child of root. |
| 15 | """ |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 16 | base_dir = os.path.normpath(os.path.abspath(base_dir)) |
| 17 | def archive_name(path): |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 18 | """Makes path relative to base_dir.""" |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 19 | path = os.path.normpath(os.path.abspath(path)) |
| 20 | common_path = os.path.commonprefix((base_dir, path)) |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 21 | archive_name = path[len(common_path):] |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 22 | if os.path.isabs(archive_name): |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 23 | archive_name = archive_name[1:] |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 24 | return os.path.join(prefix_dir, archive_name) |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 25 | def visit(tar, dirname, names): |
| 26 | for name in names: |
| 27 | path = os.path.join(dirname, name) |
| 28 | if os.path.isfile(path): |
| 29 | path_in_tar = archive_name(path) |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 30 | tar.add(path, path_in_tar) |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 31 | compression = TARGZ_DEFAULT_COMPRESSION_LEVEL |
Christopher Dunn | 9cc0bb8 | 2015-01-24 15:44:51 -0600 | [diff] [blame] | 32 | with closing(tarfile.TarFile.open(tarball_path, 'w:gz', |
| 33 | compresslevel=compression)) as tar: |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 34 | for source in sources: |
| 35 | source_path = source |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 36 | if os.path.isdir(source): |
Christopher Dunn | f357688 | 2015-01-24 16:20:25 -0600 | [diff] [blame] | 37 | for dirpath, dirnames, filenames in os.walk(source_path): |
| 38 | visit(tar, dirpath, filenames) |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 39 | else: |
| 40 | path_in_tar = archive_name(source_path) |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 41 | tar.add(source_path, path_in_tar) # filename, arcname |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 42 | |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 43 | def decompress(tarball_path, base_dir): |
Christopher Dunn | dc0f736 | 2011-06-21 21:18:49 +0000 | [diff] [blame] | 44 | """Decompress the gzipped tarball into directory base_dir. |
| 45 | """ |
Christopher Dunn | 9cc0bb8 | 2015-01-24 15:44:51 -0600 | [diff] [blame] | 46 | with closing(tarfile.TarFile.open(tarball_path)) as tar: |
Christopher Dunn | 494950a | 2015-01-24 15:29:52 -0600 | [diff] [blame] | 47 | tar.extractall(base_dir) |