Merge "Add ability to download objects to particular folder."

3 files changed, 161 insertions, 13 deletions

diff --git a/swiftclient/service.py b/swiftclient/service.py
index d7c05ca..7c55769 100644
--- a/swiftclient/service.py
+++ b/swiftclient/service.py
@@ -12,6 +12,7 @@
 # implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
 from concurrent.futures import as_completed, CancelledError, TimeoutError
 from copy import deepcopy
 from errno import EEXIST, ENOENT
@@ -162,6 +163,8 @@ _default_local_options = {
     'read_acl': None,
     'write_acl': None,
     'out_file': None,
+    'out_directory': None,
+    'remove_prefix': False,
     'no_download': False,
     'long': False,
     'totals': False,
@@ -889,7 +892,9 @@ class SwiftService(object):
                                 'no_download': False,
                                 'header': [],
                                 'skip_identical': False,
-                                'out_file': None
+                                'out_directory': None,
+                                'out_file': None,
+                                'remove_prefix': False,
                             }
 
         :returns: A generator for returning the results of the download
@@ -986,6 +991,12 @@ class SwiftService(object):
         options['skip_identical'] = (options['skip_identical'] and
                                      out_file != '-')
 
+        if options['prefix'] and options['remove_prefix']:
+            path = path[len(options['prefix']):].lstrip('/')
+
+        if options['out_directory']:
+            path = os.path.join(options['out_directory'], path)
+
         if options['skip_identical']:
             filename = out_file if out_file else path
             try:
diff --git a/swiftclient/shell.py b/swiftclient/shell.py
index 4438e9d..695a77e 100755
--- a/swiftclient/shell.py
+++ b/swiftclient/shell.py
@@ -146,9 +146,11 @@ def st_delete(parser, args, output_manager):
 
 
 st_download_options = '''[--all] [--marker] [--prefix <prefix>]
-                      [--output <out_file>] [--object-threads <threads>]
+                      [--output <out_file>] [--output-dir <out_directory>]
+                      [--object-threads <threads>]
                       [--container-threads <threads>] [--no-download]
-                      [--skip-identical] <container> <object>
+                      [--skip-identical] [--remove-prefix]
+                      <container> <object>
 '''
 
 st_download_help = '''
@@ -167,9 +169,15 @@ Optional arguments:
   --marker              Marker to use when starting a container or account
                         download.
   --prefix <prefix>     Only download items beginning with <prefix>
+  --remove-prefix       An optional flag for --prefix <prefix>, use this
+                        option to download items without <prefix>
   --output <out_file>   For a single file download, stream the output to
                         <out_file>. Specifying "-" as <out_file> will
                         redirect to stdout.
+  --output-dir <out_directory>
+                        An optional directory to which to store objects.
+                        By default, all objects are recreated in the current
+                        directory.
   --object-threads <threads>
                         Number of threads to use for downloading objects.
                         Default is 10.
@@ -204,6 +212,14 @@ def st_download(parser, args, output_manager):
         'download, stream the output to <out_file>. '
         'Specifying "-" as <out_file> will redirect to stdout.')
     parser.add_option(
+        '-D', '--output-dir', dest='out_directory',
+        help='An optional directory to which to store objects. '
+        'By default, all objects are recreated in the current directory.')
+    parser.add_option(
+        '-r', '--remove-prefix', action='store_true', dest='remove_prefix',
+        default=False, help='An optional flag for --prefix <prefix>, '
+        'use this option to download items without <prefix>.')
+    parser.add_option(
         '', '--object-threads', type=int,
         default=10, help='Number of threads to use for downloading objects. '
         'Default is 10.')
@@ -233,6 +249,12 @@ def st_download(parser, args, output_manager):
     if options.out_file and len(args) != 2:
         exit('-o option only allowed for single file downloads')
 
+    if not options.prefix:
+        options.remove_prefix = False
+
+    if options.out_directory and len(args) == 2:
+        exit('Please use -o option for single file downloads and renames')
+
     if (not args and not options.yes_all) or (args and options.yes_all):
         output_manager.error('Usage: %s download %s\n%s', BASENAME,
                              st_download_options, st_download_help)
diff --git a/tests/unit/test_service.py b/tests/unit/test_service.py
index ee1cbb6..339aca1 100644
--- a/tests/unit/test_service.py
+++ b/tests/unit/test_service.py
@@ -1029,6 +1029,17 @@ class TestServiceUpload(testtools.TestCase):
 
 class TestServiceDownload(testtools.TestCase):
 
+    def setUp(self):
+        super(TestServiceDownload, self).setUp()
+        self.opts = swiftclient.service._default_local_options.copy()
+        self.opts['no_download'] = True
+        self.obj_content = b'c' * 10
+        self.obj_etag = md5(self.obj_content).hexdigest()
+        self.obj_len = len(self.obj_content)
+
+    def _readbody(self):
+        yield self.obj_content
+
     def _assertDictEqual(self, a, b, m=None):
         # assertDictEqual is not available in py2.6 so use a shallow check
         # instead
@@ -1045,6 +1056,103 @@ class TestServiceDownload(testtools.TestCase):
                 self.assertIn(k, b, m)
                 self.assertEqual(b[k], v, m)
 
+    def test_download(self):
+        service = SwiftService()
+        with mock.patch('swiftclient.service.Connection') as mock_conn:
+            header = {'content-length': self.obj_len,
+                      'etag': self.obj_etag}
+            mock_conn.get_object.return_value = header, self._readbody()
+
+            resp = service._download_object_job(mock_conn,
+                                                'c',
+                                                'test',
+                                                self.opts)
+
+        self.assertTrue(resp['success'])
+        self.assertEqual(resp['action'], 'download_object')
+        self.assertEqual(resp['object'], 'test')
+        self.assertEqual(resp['path'], 'test')
+
+    def test_download_with_output_dir(self):
+        service = SwiftService()
+        with mock.patch('swiftclient.service.Connection') as mock_conn:
+            header = {'content-length': self.obj_len,
+                      'etag': self.obj_etag}
+            mock_conn.get_object.return_value = header, self._readbody()
+
+            options = self.opts.copy()
+            options['out_directory'] = 'temp_dir'
+            resp = service._download_object_job(mock_conn,
+                                                'c',
+                                                'example/test',
+                                                options)
+
+        self.assertTrue(resp['success'])
+        self.assertEqual(resp['action'], 'download_object')
+        self.assertEqual(resp['object'], 'example/test')
+        self.assertEqual(resp['path'], 'temp_dir/example/test')
+
+    def test_download_with_remove_prefix(self):
+        service = SwiftService()
+        with mock.patch('swiftclient.service.Connection') as mock_conn:
+            header = {'content-length': self.obj_len,
+                      'etag': self.obj_etag}
+            mock_conn.get_object.return_value = header, self._readbody()
+
+            options = self.opts.copy()
+            options['prefix'] = 'example/'
+            options['remove_prefix'] = True
+            resp = service._download_object_job(mock_conn,
+                                                'c',
+                                                'example/test',
+                                                options)
+
+        self.assertTrue(resp['success'])
+        self.assertEqual(resp['action'], 'download_object')
+        self.assertEqual(resp['object'], 'example/test')
+        self.assertEqual(resp['path'], 'test')
+
+    def test_download_with_remove_prefix_and_remove_slashes(self):
+        service = SwiftService()
+        with mock.patch('swiftclient.service.Connection') as mock_conn:
+            header = {'content-length': self.obj_len,
+                      'etag': self.obj_etag}
+            mock_conn.get_object.return_value = header, self._readbody()
+
+            options = self.opts.copy()
+            options['prefix'] = 'example'
+            options['remove_prefix'] = True
+            resp = service._download_object_job(mock_conn,
+                                                'c',
+                                                'example/test',
+                                                options)
+
+        self.assertTrue(resp['success'])
+        self.assertEqual(resp['action'], 'download_object')
+        self.assertEqual(resp['object'], 'example/test')
+        self.assertEqual(resp['path'], 'test')
+
+    def test_download_with_output_dir_and_remove_prefix(self):
+        service = SwiftService()
+        with mock.patch('swiftclient.service.Connection') as mock_conn:
+            header = {'content-length': self.obj_len,
+                      'etag': self.obj_etag}
+            mock_conn.get_object.return_value = header, self._readbody()
+
+            options = self.opts.copy()
+            options['prefix'] = 'example'
+            options['out_directory'] = 'new/dir'
+            options['remove_prefix'] = True
+            resp = service._download_object_job(mock_conn,
+                                                'c',
+                                                'example/test',
+                                                options)
+
+        self.assertTrue(resp['success'])
+        self.assertEqual(resp['action'], 'download_object')
+        self.assertEqual(resp['object'], 'example/test')
+        self.assertEqual(resp['path'], 'new/dir/test')
+
     def test_download_object_job_skip_identical(self):
         with tempfile.NamedTemporaryFile() as f:
             f.write(b'a' * 30)
@@ -1077,6 +1185,9 @@ class TestServiceDownload(testtools.TestCase):
                                        container='test_c',
                                        obj='test_o',
                                        options={'out_file': f.name,
+                                                'out_directory': None,
+                                                'prefix': None,
+                                                'remove_prefix': False,
                                                 'header': {},
                                                 'yes_all': False,
                                                 'skip_identical': True})
@@ -1129,6 +1240,9 @@ class TestServiceDownload(testtools.TestCase):
                                            container='test_c',
                                            obj='test_o',
                                            options={'out_file': f.name,
+                                                    'out_directory': None,
+                                                    'prefix': None,
+                                                    'remove_prefix': False,
                                                     'header': {},
                                                     'yes_all': False,
                                                     'skip_identical': True})
@@ -1207,6 +1321,9 @@ class TestServiceDownload(testtools.TestCase):
                                            container='test_c',
                                            obj='test_o',
                                            options={'out_file': f.name,
+                                                    'out_directory': None,
+                                                    'prefix': None,
+                                                    'remove_prefix': False,
                                                     'header': {},
                                                     'yes_all': False,
                                                     'skip_identical': True})
@@ -1268,6 +1385,9 @@ class TestServiceDownload(testtools.TestCase):
                 'auth_end_time': mock_conn.auth_end_time,
             }
 
+            options = self.opts.copy()
+            options['out_file'] = f.name
+            options['skip_identical'] = True
             s = SwiftService()
             with mock.patch('swiftclient.service.time', side_effect=range(3)):
                 with mock.patch('swiftclient.service.get_conn',
@@ -1276,11 +1396,7 @@ class TestServiceDownload(testtools.TestCase):
                         conn=mock_conn,
                         container='test_c',
                         obj='test_o',
-                        options={'out_file': f.name,
-                                 'header': {},
-                                 'no_download': True,
-                                 'yes_all': False,
-                                 'skip_identical': True})
+                        options=options)
 
             self._assertDictEqual(r, expected_r)
 
@@ -1360,6 +1476,9 @@ class TestServiceDownload(testtools.TestCase):
                 'auth_end_time': mock_conn.auth_end_time,
             }
 
+            options = self.opts.copy()
+            options['out_file'] = f.name
+            options['skip_identical'] = True
             s = SwiftService()
             with mock.patch('swiftclient.service.time', side_effect=range(3)):
                 with mock.patch('swiftclient.service.get_conn',
@@ -1368,11 +1487,7 @@ class TestServiceDownload(testtools.TestCase):
                         conn=mock_conn,
                         container='test_c',
                         obj='test_o',
-                        options={'out_file': f.name,
-                                 'header': {},
-                                 'no_download': True,
-                                 'yes_all': False,
-                                 'skip_identical': True})
+                        options=options)
 
             self._assertDictEqual(r, expected_r)
             self.assertEqual(mock_conn.get_object.mock_calls, [