1 files changed, 24 insertions, 4 deletions
diff --git a/writeexts.py b/writeexts.py
index 469b8557..af48b375 100755
--- a/writeexts.py
+++ b/writeexts.py
@@ -74,13 +74,16 @@ class WriteExtension(cliapp.Application):
         else:
             self.unmount(mp)
 
-    def get_disk_size(self):
-        '''Parse disk size from environment.'''
+    def _parse_size(self, size):
+        '''Parse a size from a string.
+        
+        Return size in bytes.
         
-        size = os.environ.get('DISK_SIZE', '1G')
+        '''
+
         m = re.match('^(\d+)([kmgKMG]?)$', size)
         if not m:
-            raise morphlib.Error('Cannot parse disk size %s' % size)
+            return None
 
         factors = {
             '': 1,
@@ -92,6 +95,23 @@ class WriteExtension(cliapp.Application):
 
         return int(m.group(1)) * factor
 
+    def _parse_size_from_environment(self, env_var, default):
+        '''Parse a size from an environment variable.'''
+
+        size = os.environ.get(env_var, default)
+        bytes = self._parse_size(size)
+        if bytes is None:
+            raise morphlib.Error('Cannot parse %s value %s' % (env_var, size))
+        return bytes
+
+    def get_disk_size(self):
+        '''Parse disk size from environment.'''
+        return self._parse_size_from_environment('DISK_SIZE', '1G')
+
+    def get_ram_size(self):
+        '''Parse RAM size from environment.'''
+        return self._parse_size_from_environment('RAM_SIZE', '1G')
+
     def create_raw_disk_image(self, filename, size):
         '''Create a raw disk image.'''