From f1bd1edb098d273fcefe16b1899789165ab01137 Mon Sep 17 00:00:00 2001
From: Minwoo Jeong <mwjeong.sci@gmail.com>
Date: Wed, 7 Jul 2021 22:43:29 +0900
Subject: [PATCH] Add method to update multiple key/values to .env

Following things were changed.
* `update_dict_to_dotenv` was added to `main.py`
* `make_env_line` was extracted from `set_key`
* `test_update_dict_to_dotenv` was added to `test_main.py`
---
 src/dotenv/__init__.py |  5 +--
 src/dotenv/main.py     | 72 ++++++++++++++++++++++++++++++++++--------
 tests/test_main.py     | 23 ++++++++++++++
 3 files changed, 84 insertions(+), 16 deletions(-)

diff --git a/src/dotenv/__init__.py b/src/dotenv/__init__.py
index 3512d101..756a410e 100644
--- a/src/dotenv/__init__.py
+++ b/src/dotenv/__init__.py
@@ -1,7 +1,7 @@
 from typing import Any, Optional
 
 from .main import (dotenv_values, find_dotenv, get_key, load_dotenv, set_key,
-                   unset_key)
+                   unset_key, update_dict_to_dotenv)
 
 
 def load_ipython_extension(ipython: Any) -> None:
@@ -46,4 +46,5 @@ def get_cli_string(
            'set_key',
            'unset_key',
            'find_dotenv',
-           'load_ipython_extension']
+           'load_ipython_extension',
+           'update_dict_to_dotenv']
diff --git a/src/dotenv/main.py b/src/dotenv/main.py
index b8d0a4e0..d365d2d4 100644
--- a/src/dotenv/main.py
+++ b/src/dotenv/main.py
@@ -135,35 +135,49 @@ def rewrite(path: Union[str, _PathLike]) -> Iterator[Tuple[IO[str], IO[str]]]:
         shutil.move(dest.name, path)
 
 
-def set_key(
-    dotenv_path: Union[str, _PathLike],
-    key_to_set: str,
-    value_to_set: str,
+def make_env_line(
+    key: str,
+    value: str,
     quote_mode: str = "always",
     export: bool = False,
-) -> Tuple[Optional[bool], str, str]:
+) -> str:
     """
-    Adds or Updates a key/value to the given .env
-
-    If the .env path given doesn't exist, fails instead of risking creating
-    an orphan .env somewhere in the filesystem
+    Make a line which format fits to .env
     """
     if quote_mode not in ("always", "auto", "never"):
         raise ValueError("Unknown quote_mode: {}".format(quote_mode))
 
     quote = (
         quote_mode == "always"
-        or (quote_mode == "auto" and not value_to_set.isalnum())
+        or (quote_mode == "auto" and not value.isalnum())
     )
 
     if quote:
-        value_out = "'{}'".format(value_to_set.replace("'", "\\'"))
+        value_out = "'{}'".format(value.replace("'", "\\'"))
     else:
-        value_out = value_to_set
+        value_out = value
     if export:
-        line_out = 'export {}={}\n'.format(key_to_set, value_out)
+        line_out = 'export {}={}\n'.format(key, value_out)
     else:
-        line_out = "{}={}\n".format(key_to_set, value_out)
+        line_out = "{}={}\n".format(key, value_out)
+
+    return line_out
+
+
+def set_key(
+    dotenv_path: Union[str, _PathLike],
+    key_to_set: str,
+    value_to_set: str,
+    quote_mode: str = "always",
+    export: bool = False,
+) -> Tuple[Optional[bool], str, str]:
+    """
+    Adds or Updates a key/value to the given .env
+
+    If the .env path given doesn't exist, fails instead of risking creating
+    an orphan .env somewhere in the filesystem
+    """
+    line_out = make_env_line(key_to_set, value_to_set, quote_mode, export)
 
     with rewrite(dotenv_path) as (source, dest):
         replaced = False
@@ -358,3 +372,33 @@ def dotenv_values(
         override=True,
         encoding=encoding,
     ).dict()
+
+
+def update_dict_to_dotenv(
+    dotenv_path: Union[str, _PathLike], 
+    env_dict: dict, 
+    quote_mode: str = "always", 
+    export: bool = False
+):
+    """
+    Adds or Updates key/value pairs in the given dictionary to the given .env
+
+    If the .env path given doesn't exist, fails instead of risking creating
+    an orphan .env somewhere in the filesystem
+    """
+    key_to_line = {}
+
+    for key_to_set, value_to_set in env_dict.items():
+        env_line = make_env_line(key_to_set, value_to_set, quote_mode, export)
+        key_to_line[key_to_set] = env_line
+
+    with rewrite(dotenv_path) as (source, dest):
+        for mapping in with_warn_for_invalid_lines(parse_stream(source)):
+            if mapping.key in key_to_line:
+                line_out = key_to_line.pop(mapping.key)
+                dest.write(line_out)
+            else:
+                dest.write(mapping.original.string)
+
+        for _, line_out in key_to_line.items():
+            dest.write(line_out)
diff --git a/tests/test_main.py b/tests/test_main.py
index 13e2791c..23ad162b 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -383,3 +383,26 @@ def test_dotenv_values_file_stream(dotenv_file):
         result = dotenv.dotenv_values(stream=f)
 
     assert result == {"a": "b"}
+
+
+@pytest.mark.parametrize(
+    "before,env_dict,after",
+    [
+        ("", {"a1": "", "a2": "b", "a3": "'b'", "a4": "\"b\""},
+         "a1=''\na2='b'\na3='\\'b\\''\na4='\"b\"'\n"),
+        ("", {"a1": "b'c", "a2": "b\"c"}, "a1='b\\'c'\na2='b\"c'\n"),
+        ("a=b\nb=c\n", {"b": "cc", "c": "d", "d": "e"},
+         "a=b\nb='cc'\nc='d'\nd='e'\n")
+    ],
+)
+def test_update_dict_to_dotenv(dotenv_file, before, env_dict, after):
+    logger = logging.getLogger("dotenv.main")
+    with open(dotenv_file, "w") as f:
+        f.write(before)
+
+    with mock.patch.object(logger, "warning") as mock_warning:
+        dotenv.update_dict_to_dotenv(dotenv_file, env_dict)
+
+    assert open(dotenv_file, "r").read() == after
+    mock_warning.assert_not_called()
+