Fix escaping of quoted values written by set_key (theskumar#236)

bbc2 · web-flow · commit b5b8ae20d0bc · 2020-02-15T08:00:56.000+05:30
* Fix PyPy jobs in CI

* Add test cases for `dotenv.set_key`

* Escape values before quoting in `dotenv.set_key`
diff --git a/.travis.yml b/.travis.yml
@@ -1,7 +1,6 @@
 language: python
 cache: pip
 os: linux
-dist: xenial
 
 jobs:
   include:
@@ -21,10 +20,8 @@ jobs:
       env: TOXENV=py38
     - python: "pypy"
       env: TOXENV=pypy
-      dist: trusty
     - python: "pypy3"
       env: TOXENV=pypy3
-      dist: trusty
 
 install:
   - pip install tox
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,10 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - Use current working directory to find `.env` when bundled by PyInstaller (#213 by
   [@gergelyk]).
 
+### Fixed
+
+- Fix escaping of quoted values written by `set_key` (#236 by [@bbc2]).
+
 ## [0.11.0] - 2020-02-07
 
 ### Added
diff --git a/src/dotenv/main.py b/src/dotenv/main.py
@@ -153,8 +153,11 @@ def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always"):
     if " " in value_to_set:
         quote_mode = "always"
 
-    line_template = '{}="{}"\n' if quote_mode == "always" else '{}={}\n'
-    line_out = line_template.format(key_to_set, value_to_set)
+    if quote_mode == "always":
+        value_out = '"{}"'.format(value_to_set.replace('"', '\\"'))
+    else:
+        value_out = value_to_set
+    line_out = "{}={}\n".format(key_to_set, value_out)
 
     with rewrite(dotenv_path) as (source, dest):
         replaced = False
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -29,14 +29,25 @@ def test_set_key_no_file(tmp_path):
     )
 
 
-def test_set_key_new(dotenv_file):
+@pytest.mark.parametrize(
+    "key,value,expected,content",
+    [
+        ("a", "", (True, "a", ""), 'a=""\n'),
+        ("a", "b", (True, "a", "b"), 'a="b"\n'),
+        ("a", "'b'", (True, "a", "b"), 'a="b"\n'),
+        ("a", "\"b\"", (True, "a", "b"), 'a="b"\n'),
+        ("a", "b'c", (True, "a", "b'c"), 'a="b\'c"\n'),
+        ("a", "b\"c", (True, "a", "b\"c"), 'a="b\\\"c"\n'),
+    ],
+)
+def test_set_key_new(dotenv_file, key, value, expected, content):
     logger = logging.getLogger("dotenv.main")
 
     with mock.patch.object(logger, "warning") as mock_warning:
-        result = dotenv.set_key(dotenv_file, "foo", "bar")
+        result = dotenv.set_key(dotenv_file, key, value)
 
-    assert result == (True, "foo", "bar")
-    assert open(dotenv_file, "r").read() == 'foo="bar"\n'
+    assert result == expected
+    assert open(dotenv_file, "r").read() == content
     mock_warning.assert_not_called()
 
 
