Merge pull request #287 from leonidlm/feature/adding-memory-limit-string-support

Adding units specification to mem_limit variable

README.md

@@ -81,6 +81,8 @@ more information on how to create port bindings and volume mappings.
 The `environment` variable accepts a dictionary or a list of strings
 in the following format `["PASSWORD=xxx"]` or `{"PASSWORD": "xxx"}`.
 
+The `mem_limit` variable accepts float values (which represent the memory limit of the created container in bytes) or a string with a units identification char ('100000b', 1000k', 128m', '1g'). If a string is specified without a units character, bytes are assumed as an intended unit.
+
 `volumes_from` and `dns` arguments raise TypeError exception if they are used
 against v1.10 of docker remote API. Those arguments should be passed to
 `start()` instead.

docker/client.py

@@ -111,6 +111,41 @@ class Client(requests.Session):
                 '{0}={1}'.format(k, v) for k, v in environment.items()
             ]
 
+        if isinstance(mem_limit, six.string_types):
+            if len(mem_limit) == 0:
+                mem_limit = 0
+            else:
+                units = {'b': 1,
+                         'k': 1024,
+                         'm': 1024*1024,
+                         'g': 1024*1024*1024}
+                suffix = mem_limit[-1].lower()
+
+                # Check if the variable is a string representation of an int
+                # without a units part. Assuming that the units are bytes.
+                if suffix.isdigit():
+                    digits_part = mem_limit
+                    suffix = 'b'
+                else:
+                    digits_part = mem_limit[:-1]
+
+                if suffix in units.keys() or suffix.isdigit():
+                    try:
+                        digits = int(digits_part)
+                    except ValueError:
+                        message = ('Failed converting the string value for'
+                                   ' mem_limit ({0}) to a number.')
+                        formatted_message = message.format(digits_part)
+                        raise errors.DockerException(formatted_message)
+
+                    mem_limit = digits * units[suffix]
+                else:
+                    message = ('The specified value for mem_limit parameter'
+                               ' ({0}) should specify the units. The postfix'
+                               ' should be one of the `b` `k` `m` `g`'
+                               ' characters')
+                    raise errors.DockerException(message.format(mem_limit))
+
         if isinstance(ports, list):
             exposed_ports = {}
             for port_definition in ports:

tests/test.py

@@ -418,6 +418,70 @@ class DockerClientTest(Cleanup, unittest.TestCase):
                          {'Content-Type': 'application/json'})
         self.assertEqual(args[1]['params'], {'name': 'marisa-kirisame'})
 
+    def test_create_container_with_mem_limit_as_int(self):
+        try:
+            self.client.create_container('busybox', 'true',
+                                         mem_limit=128.0)
+        except Exception as e:
+            self.fail('Command should not raise exception: {0}'.format(e))
+
+        args = fake_request.call_args
+        data = json.loads(args[1]['data'])
+        self.assertEqual(data['Memory'], 128.0)
+
+    def test_create_container_with_mem_limit_as_string(self):
+        try:
+            self.client.create_container('busybox', 'true',
+                                         mem_limit='128')
+        except Exception as e:
+            self.fail('Command should not raise exception: {0}'.format(e))
+
+        args = fake_request.call_args
+        data = json.loads(args[1]['data'])
+        self.assertEqual(data['Memory'], 128.0)
+
+    def test_create_container_with_mem_limit_as_string_with_k_unit(self):
+        try:
+            self.client.create_container('busybox', 'true',
+                                         mem_limit='128k')
+        except Exception as e:
+            self.fail('Command should not raise exception: {0}'.format(e))
+
+        args = fake_request.call_args
+        data = json.loads(args[1]['data'])
+        self.assertEqual(data['Memory'], 128.0 * 1024)
+
+    def test_create_container_with_mem_limit_as_string_with_m_unit(self):
+        try:
+            self.client.create_container('busybox', 'true',
+                                         mem_limit='128m')
+        except Exception as e:
+            self.fail('Command should not raise exception: {0}'.format(e))
+
+        args = fake_request.call_args
+        data = json.loads(args[1]['data'])
+        self.assertEqual(data['Memory'], 128.0 * 1024 * 1024)
+
+    def test_create_container_with_mem_limit_as_string_with_g_unit(self):
+        try:
+            self.client.create_container('busybox', 'true',
+                                         mem_limit='128g')
+        except Exception as e:
+            self.fail('Command should not raise exception: {0}'.format(e))
+
+        args = fake_request.call_args
+        data = json.loads(args[1]['data'])
+        self.assertEqual(data['Memory'], 128.0 * 1024 * 1024 * 1024)
+
+    def test_create_container_with_mem_limit_as_string_with_wrong_value(self):
+        self.assertRaises(docker.errors.DockerException,
+                          self.client.create_container,
+                          'busybox', 'true', mem_limit='128p')
+
+        self.assertRaises(docker.errors.DockerException,
+                          self.client.create_container,
+                          'busybox', 'true', mem_limit='1f28')
+
     def test_start_container(self):
         try:
             self.client.start(fake_api.FAKE_CONTAINER_ID)