Support for configurable quoting mode, references theskumar#15

Hugo Chinchilla Carbonell · Hugo Chinchilla Carbonell · commit e88ce39e4fdb · 2016-03-15T13:11:14.000+01:00
diff --git a/Makefile b/Makefile
@@ -1,4 +1,4 @@
-.PHONY: clean-pyc clean-build
+.PHONY: clean-pyc clean-build test
 
 clean: clean-build clean-pyc
 
diff --git a/dotenv/cli.py b/dotenv/cli.py
@@ -9,11 +9,15 @@
 @click.option('-f', '--file', default=os.path.join(os.getcwd(), '.env'),
               type=click.Path(exists=True),
               help="Location of the .env file, defaults to .env file in current working directory.")
+@click.option('-q', '--quote', default='always',
+              type=click.Choice(['always', 'never', 'auto']),
+              help="Whether to quote or not the variable values. Default mode is always.")
 @click.pass_context
-def cli(ctx, file):
+def cli(ctx, file, quote):
     '''This script is used to set, get or unset values from a .env file.'''
     ctx.obj = {}
     ctx.obj['FILE'] = file
+    ctx.obj['QUOTE'] = quote
 
     # Need to investigate if this can actually work or if the scope of the new environ variables
     # Expires when python exits
@@ -43,7 +47,8 @@ def list(ctx):
 def set(ctx, key, value):
     '''Store the given key/value.'''
     file = ctx.obj['FILE']
-    success, key, value = set_key(file, key, value)
+    quote = ctx.obj['QUOTE']
+    success, key, value = set_key(file, key, value, quote)
     if success:
         click.echo('%s="%s"' % (key, value))
     else:
@@ -69,7 +74,8 @@ def get(ctx, key):
 def unset(ctx, key):
     '''Removes the given key.'''
     file = ctx.obj['FILE']
-    success, key = unset_key(file, key)
+    quote = ctx.obj['QUOTE']
+    success, key = unset_key(file, key, quote)
     if success:
         click.echo("Successfully removed %s" % key)
     else:
diff --git a/dotenv/main.py b/dotenv/main.py
@@ -36,7 +36,7 @@ def get_key(dotenv_path, key_to_get):
         return None
 
 
-def set_key(dotenv_path, key_to_set, value_to_set):
+def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always"):
     """
     Adds or Updates a key/value to the given .env
 
@@ -50,11 +50,11 @@ def set_key(dotenv_path, key_to_set, value_to_set):
         return None, key_to_set, value_to_set
     dotenv_as_dict = OrderedDict(parse_dotenv(dotenv_path))
     dotenv_as_dict[key_to_set] = value_to_set
-    success = flatten_and_write(dotenv_path, dotenv_as_dict)
+    success = flatten_and_write(dotenv_path, dotenv_as_dict, quote_mode)
     return success, key_to_set, value_to_set
 
 
-def unset_key(dotenv_path, key_to_unset):
+def unset_key(dotenv_path, key_to_unset, quote_mode="always"):
     """
     Removes a given key from the given .env
 
@@ -71,7 +71,7 @@ def unset_key(dotenv_path, key_to_unset):
     else:
         warnings.warn("key %s not removed from %s - key doesn't exist." % (key_to_unset, dotenv_path))
         return None, key_to_unset
-    success = flatten_and_write(dotenv_path, dotenv_as_dict)
+    success = flatten_and_write(dotenv_path, dotenv_as_dict, quote_mode)
     return success, key_to_unset
 
 
@@ -86,8 +86,12 @@ def parse_dotenv(dotenv_path):
             yield k, v
 
 
-def flatten_and_write(dotenv_path, dotenv_as_dict):
+def flatten_and_write(dotenv_path, dotenv_as_dict, quote_mode="always"):
     with open(dotenv_path, "w") as f:
         for k, v in dotenv_as_dict.items():
-            f.write('%s="%s"\n' % (k, v))
+            _mode = quote_mode
+            if _mode == "auto" and " " in v:
+                _mode = "always"
+            str_format = '%s="%s"\n' if _mode == "always" else '%s=%s\n'
+            f.write(str_format % (k, v))
     return True
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -35,12 +35,21 @@ def test_unset():
 
 
 def test_console_script(cli):
+    TEST_COMBINATIONS = (
+        # quote_mode, var_name, var_value, expected_result
+        ("always", "HELLO", "WORLD", 'HELLO="WORLD"\n'),
+        ("never", "HELLO", "WORLD", 'HELLO=WORLD\n'),
+        ("auto", "HELLO", "WORLD", 'HELLO=WORLD\n'),
+        ("auto", "HELLO", "HELLO WORLD", 'HELLO="HELLO WORLD"\n'),
+    )
     with cli.isolated_filesystem():
-        sh.touch(dotenv_path)
-        sh.dotenv('-f', dotenv_path, 'set', 'HELLO', 'WORLD')
-        output = sh.dotenv('-f', dotenv_path, 'get', 'HELLO', )
-        assert output == 'HELLO="WORLD"\n'
-        sh.rm(dotenv_path)
+        for quote_mode,variable,value,expected_result in TEST_COMBINATIONS:
+            sh.touch(dotenv_path)
+            sh.dotenv('-f', dotenv_path, '-q', quote_mode, 'set', variable, value)
+            output = sh.cat(dotenv_path)
+            assert output == expected_result
+            sh.rm(dotenv_path)
+
 
     # should fail for not existing file
     result = cli.invoke(dotenv.cli.set, ['my_key', 'my_value'])

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-.PHONY: clean-pyc clean-build`
	`1`	`+.PHONY: clean-pyc clean-build test`
`2`	`2`
`3`	`3`	`clean: clean-build clean-pyc`
`4`	`4`