Skip to content

Commit f1bd1ed

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 36516a7 commit f1bd1ed

File tree

3 files changed

+84
-16
lines changed

3 files changed

+84
-16
lines changed

src/dotenv/__init__.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from typing import Any, Optional
22

33
from .main import (dotenv_values, find_dotenv, get_key, load_dotenv, set_key,
4-
unset_key)
4+
unset_key, update_dict_to_dotenv)
55

66

77
def load_ipython_extension(ipython: Any) -> None:
@@ -46,4 +46,5 @@ def get_cli_string(
4646
'set_key',
4747
'unset_key',
4848
'find_dotenv',
49-
'load_ipython_extension']
49+
'load_ipython_extension',
50+
'update_dict_to_dotenv']

src/dotenv/main.py

Lines changed: 58 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -135,35 +135,49 @@ def rewrite(path: Union[str, _PathLike]) -> Iterator[Tuple[IO[str], IO[str]]]:
135135
shutil.move(dest.name, path)
136136

137137

138-
def set_key(
139-
dotenv_path: Union[str, _PathLike],
140-
key_to_set: str,
141-
value_to_set: str,
138+
def make_env_line(
139+
key: str,
140+
value: str,
142141
quote_mode: str = "always",
143142
export: bool = False,
144-
) -> Tuple[Optional[bool], str, str]:
143+
) -> str:
145144
"""
146-
Adds or Updates a key/value to the given .env
147-
148-
If the .env path given doesn't exist, fails instead of risking creating
149-
an orphan .env somewhere in the filesystem
145+
Make a line which format fits to .env
150146
"""
151147
if quote_mode not in ("always", "auto", "never"):
152148
raise ValueError("Unknown quote_mode: {}".format(quote_mode))
153149

154150
quote = (
155151
quote_mode == "always"
156-
or (quote_mode == "auto" and not value_to_set.isalnum())
152+
or (quote_mode == "auto" and not value.isalnum())
157153
)
158154

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

168182
with rewrite(dotenv_path) as (source, dest):
169183
replaced = False
@@ -358,3 +372,33 @@ def dotenv_values(
358372
override=True,
359373
encoding=encoding,
360374
).dict()
375+
376+
377+
def update_dict_to_dotenv(
378+
dotenv_path: Union[str, _PathLike],
379+
env_dict: dict,
380+
quote_mode: str = "always",
381+
export: bool = False
382+
):
383+
"""
384+
Adds or Updates key/value pairs in the given dictionary to the given .env
385+
386+
If the .env path given doesn't exist, fails instead of risking creating
387+
an orphan .env somewhere in the filesystem
388+
"""
389+
key_to_line = {}
390+
391+
for key_to_set, value_to_set in env_dict.items():
392+
env_line = make_env_line(key_to_set, value_to_set, quote_mode, export)
393+
key_to_line[key_to_set] = env_line
394+
395+
with rewrite(dotenv_path) as (source, dest):
396+
for mapping in with_warn_for_invalid_lines(parse_stream(source)):
397+
if mapping.key in key_to_line:
398+
line_out = key_to_line.pop(mapping.key)
399+
dest.write(line_out)
400+
else:
401+
dest.write(mapping.original.string)
402+
403+
for _, line_out in key_to_line.items():
404+
dest.write(line_out)

tests/test_main.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -383,3 +383,26 @@ def test_dotenv_values_file_stream(dotenv_file):
383383
result = dotenv.dotenv_values(stream=f)
384384

385385
assert result == {"a": "b"}
386+
387+
388+
@pytest.mark.parametrize(
389+
"before,env_dict,after",
390+
[
391+
("", {"a1": "", "a2": "b", "a3": "'b'", "a4": "\"b\""},
392+
"a1=''\na2='b'\na3='\\'b\\''\na4='\"b\"'\n"),
393+
("", {"a1": "b'c", "a2": "b\"c"}, "a1='b\\'c'\na2='b\"c'\n"),
394+
("a=b\nb=c\n", {"b": "cc", "c": "d", "d": "e"},
395+
"a=b\nb='cc'\nc='d'\nd='e'\n")
396+
],
397+
)
398+
def test_update_dict_to_dotenv(dotenv_file, before, env_dict, after):
399+
logger = logging.getLogger("dotenv.main")
400+
with open(dotenv_file, "w") as f:
401+
f.write(before)
402+
403+
with mock.patch.object(logger, "warning") as mock_warning:
404+
dotenv.update_dict_to_dotenv(dotenv_file, env_dict)
405+
406+
assert open(dotenv_file, "r").read() == after
407+
mock_warning.assert_not_called()
408+

0 commit comments

Comments
 (0)