From 8d54604bf4cf215573b521cd29d48cffd8b0bbbc Mon Sep 17 00:00:00 2001
From: Bryan Ricker <978899+bricker@users.noreply.github.com>
Date: Sun, 19 Mar 2023 13:16:29 -0700
Subject: [PATCH] feat: accept None to 'set_key' to write value-less
 environment variable

---
 src/dotenv/main.py | 34 ++++++++++++++++++++--------------
 tests/test_main.py |  1 +
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/src/dotenv/main.py b/src/dotenv/main.py
index f40c20ea..908eaf66 100644
--- a/src/dotenv/main.py
+++ b/src/dotenv/main.py
@@ -147,11 +147,11 @@ def rewrite(
 def set_key(
     dotenv_path: StrPath,
     key_to_set: str,
-    value_to_set: str,
+    value_to_set: Optional[str],
     quote_mode: str = "always",
     export: bool = False,
     encoding: Optional[str] = "utf-8",
-) -> Tuple[Optional[bool], str, str]:
+) -> Tuple[Optional[bool], str, Optional[str]]:
     """
     Adds or Updates a key/value to the given .env
 
@@ -161,19 +161,25 @@ def set_key(
     if quote_mode not in ("always", "auto", "never"):
         raise ValueError(f"Unknown quote_mode: {quote_mode}")
 
-    quote = (
-        quote_mode == "always"
-        or (quote_mode == "auto" and not value_to_set.isalnum())
-    )
-
-    if quote:
-        value_out = "'{}'".format(value_to_set.replace("'", "\\'"))
-    else:
-        value_out = value_to_set
-    if export:
-        line_out = f'export {key_to_set}={value_out}\n'
+    if value_to_set is None:
+        if export:
+            line_out = f'export {key_to_set}\n'
+        else:
+            line_out = f"{key_to_set}\n"
     else:
-        line_out = f"{key_to_set}={value_out}\n"
+        quote = (
+            quote_mode == "always"
+            or (quote_mode == "auto" and not value_to_set.isalnum())
+        )
+
+        if quote:
+            value_out = "'{}'".format(value_to_set.replace("'", "\\'"))
+        else:
+            value_out = value_to_set
+        if export:
+            line_out = f'export {key_to_set}={value_out}\n'
+        else:
+            line_out = f"{key_to_set}={value_out}\n"
 
     with rewrite(dotenv_path, encoding=encoding) as (source, dest):
         replaced = False
diff --git a/tests/test_main.py b/tests/test_main.py
index 9c895851..308676a2 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -26,6 +26,7 @@ def test_set_key_no_file(tmp_path):
     "before,key,value,expected,after",
     [
         ("", "a", "", (True, "a", ""), "a=''\n"),
+        ("", "a", None, (True, "a", None), "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"),