diff -r 27edae50e62c Lib/test/test_zipfile.py --- a/Lib/test/test_zipfile.py Sun Oct 09 20:19:21 2016 -0700 +++ b/Lib/test/test_zipfile.py Sun Oct 16 18:27:41 2016 +0300 @@ -12,7 +12,8 @@ import unittest from tempfile import TemporaryFile from random import randint, random, getrandbits -from test.support import (TESTFN, findfile, unlink, rmtree, +from test.support import script_helper +from test.support import (TESTFN, findfile, unlink, rmtree, temp_dir, requires_zlib, requires_bz2, requires_lzma, captured_stdout, check_warnings) @@ -2041,5 +2042,74 @@ class ZipInfoTests(unittest.TestCase): self.assertEqual(zi.compress_type, zipfile.ZIP_STORED) self.assertEqual(zi.file_size, 0) + +class CommandLineTest(unittest.TestCase): + + def zipfilecmd(self, *args, **kwargs): + rc, out, err = script_helper.assert_python_ok('-m', 'zipfile', *args, + **kwargs) + return out.replace(os.linesep.encode(), b'\n') + + def zipfilecmd_failure(self, *args): + return script_helper.assert_python_failure('-m', 'zipfile', *args) + + def test_test_command(self): + zip_name = findfile('zipdir.zip') + for opt in '-t', '--test': + out = self.zipfilecmd(opt, zip_name) + self.assertEqual(out.rstrip(), b'Done testing') + zip_name = findfile('testtar.tar') + rc, out, err = self.zipfilecmd_failure('-t', zip_name) + self.assertEqual(out, b'') + + def test_list_command(self): + zip_name = findfile('zipdir.zip') + t = io.StringIO() + with zipfile.ZipFile(zip_name, 'r') as tf: + tf.printdir(t) + expected = t.getvalue().encode('ascii', 'backslashreplace') + for opt in '-l', '--list': + out = self.zipfilecmd(opt, zip_name, + PYTHONIOENCODING='ascii:backslashreplace') + self.assertEqual(out, expected) + + def test_create_command(self): + self.addCleanup(unlink, TESTFN) + with open(TESTFN, 'w') as f: + f.write('test 1') + os.mkdir(TESTFNDIR) + self.addCleanup(rmtree, TESTFNDIR) + with open(os.path.join(TESTFNDIR, 'file.txt'), 'w') as f: + f.write('test 2') + files = [TESTFN, TESTFNDIR] + namelist = [TESTFN, TESTFNDIR + '/', TESTFNDIR + '/file.txt'] + for opt in '-c', '--create': + try: + out = self.zipfilecmd(opt, TESTFN2, *files) + self.assertEqual(out, b'') + with zipfile.ZipFile(TESTFN2) as zf: + self.assertEqual(zf.namelist(), namelist) + self.assertEqual(zf.read(namelist[0]), b'test 1') + self.assertEqual(zf.read(namelist[2]), b'test 2') + finally: + unlink(TESTFN2) + + def test_extract_command(self): + zip_name = findfile('zipdir.zip') + for opt in '-e', '--extract': + with temp_dir() as extdir: + out = self.zipfilecmd(opt, zip_name, extdir) + self.assertEqual(out, b'') + with zipfile.ZipFile(zip_name) as zf: + for zi in zf.infolist(): + path = os.path.join(extdir, + zi.filename.replace('/', os.sep)) + if zi.is_dir(): + self.assertTrue(os.path.isdir(path)) + else: + self.assertTrue(os.path.isfile(path)) + with open(path, 'rb') as f: + self.assertEqual(f.read(), zf.read(zi)) + if __name__ == "__main__": unittest.main() diff -r 27edae50e62c Lib/zipfile.py --- a/Lib/zipfile.py Sun Oct 09 20:19:21 2016 -0700 +++ b/Lib/zipfile.py Sun Oct 16 18:27:41 2016 +0300 @@ -1950,51 +1950,45 @@ class PyZipFile(ZipFile): return (fname, archivename) -def main(args = None): - import textwrap - USAGE=textwrap.dedent("""\ - Usage: - zipfile.py -l zipfile.zip # Show listing of a zipfile - zipfile.py -t zipfile.zip # Test if a zipfile is valid - zipfile.py -e zipfile.zip target # Extract zipfile into target dir - zipfile.py -c zipfile.zip src ... # Create zipfile from sources - """) - if args is None: - args = sys.argv[1:] +def main(args=None): + import argparse - if not args or args[0] not in ('-l', '-c', '-e', '-t'): - print(USAGE, file=sys.stderr) - sys.exit(1) + description = 'A simple command line interface for zipfile module.' + parser = argparse.ArgumentParser(description=description) + group = parser.add_mutually_exclusive_group() + group.add_argument('-l', '--list', metavar='', + help='Show listing of a zipfile') + group.add_argument('-e', '--extract', nargs=2, + metavar=('', ''), + help='Extract zipfile into target dir') + group.add_argument('-c', '--create', nargs='+', + metavar=('', ''), + help='Create zipfile from sources') + group.add_argument('-t', '--test', metavar='', + help='Test if a zipfile is valid') + args = parser.parse_args(args) - if args[0] == '-l': - if len(args) != 2: - print(USAGE, file=sys.stderr) - sys.exit(1) - with ZipFile(args[1], 'r') as zf: - zf.printdir() - - elif args[0] == '-t': - if len(args) != 2: - print(USAGE, file=sys.stderr) - sys.exit(1) - with ZipFile(args[1], 'r') as zf: + if args.test is not None: + src = args.test + with ZipFile(src, 'r') as zf: badfile = zf.testzip() if badfile: print("The following enclosed file is corrupted: {!r}".format(badfile)) print("Done testing") - elif args[0] == '-e': - if len(args) != 3: - print(USAGE, file=sys.stderr) - sys.exit(1) + elif args.list is not None: + src = args.list + with ZipFile(src, 'r') as zf: + zf.printdir() - with ZipFile(args[1], 'r') as zf: - zf.extractall(args[2]) + elif args.extract is not None: + src, curdir = args.extract + with ZipFile(src, 'r') as zf: + zf.extractall(curdir) - elif args[0] == '-c': - if len(args) < 3: - print(USAGE, file=sys.stderr) - sys.exit(1) + elif args.create is not None: + zip_name = args.create.pop(0) + files = args.create def addToZip(zf, path, zippath): if os.path.isfile(path): @@ -2007,8 +2001,8 @@ def main(args = None): os.path.join(path, nm), os.path.join(zippath, nm)) # else: ignore - with ZipFile(args[1], 'w') as zf: - for path in args[2:]: + with ZipFile(zip_name, 'w') as zf: + for path in files: zippath = os.path.basename(path) if not zippath: zippath = os.path.basename(os.path.dirname(path)) @@ -2016,5 +2010,8 @@ def main(args = None): zippath = '' addToZip(zf, path, zippath) + else: + parser.exit(2, parser.format_usage()) + if __name__ == "__main__": main()