Add function to write header block to .env file

Author: ankona

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)
+from .main import (dotenv_values, find_dotenv, get_key, load_dotenv, set_header,
+                   set_key, unset_key)
 
 
 def load_ipython_extension(ipython: Any) -> None:
@@ -42,6 +42,7 @@ def get_cli_string(
 __all__ = ['get_cli_string',
            'load_dotenv',
            'dotenv_values',
+           'set_header',
            'get_key',
            'set_key',
            'unset_key',

src/dotenv/main.py

@@ -396,3 +396,36 @@ def dotenv_values(
         override=True,
         encoding=encoding,
     ).dict()
+
+
+def set_header(
+    dotenv_path: StrPath,
+    header: str,
+    encoding: Optional[str] = "utf-8",
+) -> Tuple[bool, Optional[str]]:
+    """
+    Adds or Updates a header in the .env file
+
+    Parameters:
+        dotenv_path: Absolute or relative path to .env file.
+        header: The desired header block
+        encoding: Encoding to be used to read the file.
+    Returns:
+        Bool: True if at least one environment variable is set else False
+        Str: The header that was written
+    """
+    with rewrite(dotenv_path, encoding=encoding) as (source, dest):
+        if not header or not header.strip():
+            logger.info("Ignoring empty header.")
+            return False, header
+
+        header = header.strip()
+        lines = header.split("\n")
+        for i, line in enumerate(lines):
+            if not line.startswith("# "):
+                lines[i] = f"# {line}"
+        header = "\n".join(lines)
+        text = "".join(atom for atom in source.readlines() if not atom.startswith("#"))
+        dest.write(f"{header}\n{text}\n")
+
+    return True, header

tests/test_main.py

@@ -396,3 +396,54 @@ def test_dotenv_values_file_stream(dotenv_path):
         result = dotenv.dotenv_values(stream=f)
 
     assert result == {"a": "b"}
+
+
+@pytest.mark.parametrize(
+    "header",
+    [
+        "",
+        " ",
+        None,
+    ],
+)
+def test_set_header_empty(dotenv_path, header):
+    logger = logging.getLogger("dotenv.main")
+
+    with mock.patch.object(logger, "info") as mock_info:
+        result, *_ = dotenv.set_header(dotenv_path, header)
+        assert not result
+
+    mock_info.assert_called()
+
+
+@pytest.mark.parametrize(
+    "new_header, expected, expected_header, old_header, content",
+    [
+        ("# header only", True, "# header only", "", ""),
+        ("# single-line header", True, "# single-line header", "", "a=b\nc=d"),
+        ("# multi-line\n# header", True, "# multi-line\n# header", "", "a=b\nc=d"),
+        (
+            "Single-line, no comment header",
+            True,
+            "# Single-line, no comment header",
+            "",
+            "a=b",
+        ),
+        ("# New header", True, "# New header", "# Old header", "a=b\nc=d"),
+    ],
+)
+def test_set_header(
+    dotenv_path, new_header, expected, expected_header, old_header, content
+):
+    logger = logging.getLogger("dotenv.main")
+    dotenv_path.write_text(f"{old_header}\n{content}")
+
+    with mock.patch.object(logger, "warning") as mock_warning:
+        result, written = dotenv.set_header(dotenv_path, new_header)
+        assert result == expected
+        assert written == expected_header
+
+    text = dotenv_path.read_text()
+    assert content in text
+    assert expected_header in text
+    mock_warning.assert_not_called()