Merge "Allow segment size to be specified in a human readable way."

swiftclient/shell.py

@@ -689,7 +689,9 @@ def st_upload(parser, args, output_manager):
         '-S', '--segment-size', dest='segment_size', help='Upload files '
         'in segments no larger than <size> (in Bytes) and then create a '
         '"manifest" file that will download all the segments as if it were '
-        'the original file.')
+        'the original file. Sizes may also be expressed as bytes with the '
+        'B suffix, kilobytes with the K suffix, megabytes with the M suffix '
+        'or gigabytes with the G suffix.')
     parser.add_option(
         '-C', '--segment-container', dest='segment_container',
         help='Upload the segments into the specified container. '
@@ -741,6 +743,20 @@ def st_upload(parser, args, output_manager):
         else:
             orig_path = files[0]
 
+    if options.segment_size:
+        try:
+            # If segment size only has digits assume it is bytes
+            int(options.segment_size)
+        except ValueError:
+            try:
+                size_mod = "BKMG".index(options.segment_size[-1].upper())
+                multiplier = int(options.segment_size[:-1])
+            except ValueError:
+                output_manager.error("Invalid segment size")
+                return
+
+            options.segment_size = str((1024 ** size_mod) * multiplier)
+
     _opts = vars(options)
     _opts['object_uu_threads'] = options.object_threads
     with SwiftService(options=_opts) as swift:

tests/unit/test_shell.py

@@ -444,6 +444,56 @@ class TestShell(unittest.TestCase):
         swiftclient.shell.main(argv)
         connection.return_value.get_capabilities.assert_called_with(None)
 
+    def test_human_readable_upload_segment_size(self):
+        def _check_expected(x, expected):
+            actual = x.call_args_list[-1][1]["options"]["segment_size"]
+            self.assertEqual(int(actual), expected)
+
+        mock_out = mock.MagicMock(spec=swiftclient.shell.OutputManager)
+        mock_out.__enter__.return_value = mock_out
+        mock_out.return_value = mock_out
+        type(mock_out).error_count = mock.PropertyMock(return_value=0)
+
+        mock_swift = mock.MagicMock(spec=swiftclient.shell.SwiftService)
+
+        with mock.patch("swiftclient.shell.SwiftService", mock_swift):
+            with mock.patch('swiftclient.shell.OutputManager', mock_out):
+                # Test new behaviour with both upper and lower case
+                # trailing characters
+                argv = ["", "upload", "-S", "1B", "container", "object"]
+                swiftclient.shell.main(argv)
+                _check_expected(mock_swift, 1)
+
+                argv = ["", "upload", "-S", "1K", "container", "object"]
+                swiftclient.shell.main(argv)
+                _check_expected(mock_swift, 1024)
+
+                argv = ["", "upload", "-S", "1m", "container", "object"]
+                swiftclient.shell.main(argv)
+                _check_expected(mock_swift, 1048576)
+
+                argv = ["", "upload", "-S", "1G", "container", "object"]
+                swiftclient.shell.main(argv)
+                _check_expected(mock_swift, 1073741824)
+
+                # Test old behaviour is not affected
+                argv = ["", "upload", "-S", "12345", "container", "object"]
+                swiftclient.shell.main(argv)
+                _check_expected(mock_swift, 12345)
+
+                # Test invalid states
+                argv = ["", "upload", "-S", "1234X", "container", "object"]
+                swiftclient.shell.main(argv)
+                mock_out.error.assert_called_with("Invalid segment size")
+
+                argv = ["", "upload", "-S", "K1234", "container", "object"]
+                swiftclient.shell.main(argv)
+                mock_out.error.assert_called_with("Invalid segment size")
+
+                argv = ["", "upload", "-S", "K", "container", "object"]
+                swiftclient.shell.main(argv)
+                mock_out.error.assert_called_with("Invalid segment size")
+
 
 class TestSubcommandHelp(unittest.TestCase):