Skip to content

Commit 3c9bfe7

Browse files
committed
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`
1 parent 5c7f43f commit 3c9bfe7

File tree

3 files changed

+72
-13
lines changed

3 files changed

+72
-13
lines changed

src/dotenv/__init__.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from .compat import IS_TYPE_CHECKING
2-
from .main import load_dotenv, get_key, set_key, unset_key, find_dotenv, dotenv_values
2+
from .main import load_dotenv, get_key, set_key, unset_key, find_dotenv, dotenv_values, update_dict_to_dotenv
33

44
if IS_TYPE_CHECKING:
55
from typing import Any, Optional
@@ -43,4 +43,5 @@ def get_cli_string(path=None, action=None, key=None, value=None, quote=None):
4343
'set_key',
4444
'unset_key',
4545
'find_dotenv',
46-
'load_ipython_extension']
46+
'load_ipython_extension',
47+
'update_dict_to_dotenv']

src/dotenv/main.py

Lines changed: 47 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -143,30 +143,40 @@ def rewrite(path):
143143
shutil.move(dest.name, path)
144144

145145

146-
def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always", export=False):
147-
# type: (_PathLike, Text, Text, Text, bool) -> Tuple[Optional[bool], Text, Text]
146+
def make_env_line(key, value, quote_mode="always", export=False):
147+
# type: (Text, Text, Text, bool) -> Text
148148
"""
149-
Adds or Updates a key/value to the given .env
150-
151-
If the .env path given doesn't exist, fails instead of risking creating
152-
an orphan .env somewhere in the filesystem
149+
Make a line which format fits to .env
153150
"""
154151
if quote_mode not in ("always", "auto", "never"):
155152
raise ValueError("Unknown quote_mode: {}".format(quote_mode))
156153

157154
quote = (
158155
quote_mode == "always"
159-
or (quote_mode == "auto" and not value_to_set.isalnum())
156+
or (quote_mode == "auto" and not value.isalnum())
160157
)
161158

162159
if quote:
163-
value_out = "'{}'".format(value_to_set.replace("'", "\\'"))
160+
value_out = "'{}'".format(value.replace("'", "\\'"))
164161
else:
165-
value_out = value_to_set
162+
value_out = value
166163
if export:
167-
line_out = 'export {}={}\n'.format(key_to_set, value_out)
164+
line_out = 'export {}={}\n'.format(key, value_out)
168165
else:
169-
line_out = "{}={}\n".format(key_to_set, value_out)
166+
line_out = "{}={}\n".format(key, value_out)
167+
168+
return line_out
169+
170+
171+
def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always", export=False):
172+
# type: (_PathLike, Text, Text, Text, bool) -> Tuple[Optional[bool], Text, Text]
173+
"""
174+
Adds or Updates a key/value to the given .env
175+
176+
If the .env path given doesn't exist, fails instead of risking creating
177+
an orphan .env somewhere in the filesystem
178+
"""
179+
line_out = make_env_line(key_to_set, value_to_set, quote_mode, export)
170180

171181
with rewrite(dotenv_path) as (source, dest):
172182
replaced = False
@@ -356,3 +366,29 @@ def dotenv_values(
356366
override=True,
357367
encoding=encoding,
358368
).dict()
369+
370+
371+
def update_dict_to_dotenv(dotenv_path, env_dict, quote_mode="always", export=False):
372+
# type: (_PathLike, Dict[Text, Optional[Text]], Text, bool) -> None
373+
"""
374+
Adds or Updates key/value pairs in the given dictionary to the given .env
375+
376+
If the .env path given doesn't exist, fails instead of risking creating
377+
an orphan .env somewhere in the filesystem
378+
"""
379+
key_to_line = {}
380+
381+
for key_to_set, value_to_set in env_dict.items():
382+
env_line = make_env_line(key_to_set, value_to_set, quote_mode, export)
383+
key_to_line[key_to_set] = env_line
384+
385+
with rewrite(dotenv_path) as (source, dest):
386+
for mapping in with_warn_for_invalid_lines(parse_stream(source)):
387+
if mapping.key in key_to_line:
388+
line_out = key_to_line.pop(mapping.key)
389+
dest.write(line_out)
390+
else:
391+
dest.write(mapping.original.string)
392+
393+
for _, line_out in key_to_line.items():
394+
dest.write(line_out)

tests/test_main.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,3 +367,25 @@ def test_dotenv_values_stream(env, string, interpolate, expected):
367367
result = dotenv.dotenv_values(stream=stream, interpolate=interpolate)
368368

369369
assert result == expected
370+
371+
372+
@pytest.mark.parametrize(
373+
"before,env_dict,after",
374+
[
375+
("", {"a1": "", "a2": "b", "a3": "'b'", "a4": "\"b\""},
376+
"a1=''\na2='b'\na3='\\'b\\''\na4='\"b\"'\n"),
377+
("", {"a1": "b'c", "a2": "b\"c"}, "a1='b\\'c'\na2='b\"c'\n"),
378+
("a=b\nb=c\n", {"b": "cc", "c": "d", "d": "e"},
379+
"a=b\nb='cc'\nc='d'\nd='e'\n")
380+
],
381+
)
382+
def test_update_dict_to_dotenv(dotenv_file, before, env_dict, after):
383+
logger = logging.getLogger("dotenv.main")
384+
with open(dotenv_file, "w") as f:
385+
f.write(before)
386+
387+
with mock.patch.object(logger, "warning") as mock_warning:
388+
dotenv.update_dict_to_dotenv(dotenv_file, env_dict)
389+
390+
assert open(dotenv_file, "r").read() == after
391+
mock_warning.assert_not_called()

0 commit comments

Comments
 (0)