Merge pull request #6 from hija/feature/keep-values

Feature/keep values

clean_dotenv/_main.py

@@ -1,10 +1,10 @@
 import os
 import argparse
-from typing import Iterator
+from typing import Iterator, List
 import clean_dotenv._parser as DotEnvParser
 
 
-def _clean_env(path_to_env: str):
+def _clean_env(path_to_env: str, values_to_keep: List[str] = []):
     # Open the .env file and remove the sensitive data
     # We rely on python-dotenv to parse the file, since we do not want to write our own parser
     dotenv_elements = DotEnvParser.parse_stream(open(path_to_env))
@@ -22,7 +22,11 @@ def _clean_env(path_to_env: str):
                 print(dotenv_element.export, end="", file=example_env_f)
             if dotenv_element.key:
                 print(
-                    f"{dotenv_element.key}={dotenv_element.separator}{dotenv_element.separator}",
+                    (
+                        f"{dotenv_element.key}={dotenv_element.separator}{dotenv_element.value}{dotenv_element.separator}"
+                        if dotenv_element.key in values_to_keep
+                        else f"{dotenv_element.key}={dotenv_element.separator}{dotenv_element.separator}"
+                    ),
                     end="",
                     file=example_env_f,
                 )
@@ -40,11 +44,11 @@ def _find_dotenv_files(path_to_root: str) -> Iterator[str]:
             yield entry.path
 
 
-def _main(path_to_root: str):
+def _main(path_to_root: str, values_to_keep: List[str] = []):
     # Find possible .env files
     for dotenv_file in _find_dotenv_files(path_to_root):
         # Clean dotenv file
-        _clean_env(dotenv_file)
+        _clean_env(path_to_env=dotenv_file, values_to_keep=values_to_keep)
 
 
 def main():
@@ -57,8 +61,16 @@ def main():
         help="Root path in which .env files shall be looked for",
         default=os.getcwd(),
     )
+    parser.add_argument(
+        "-k",
+        "--keep",
+        nargs="*",
+        help="Variables which shall not be cleaned by clean-dotenv. Separate values by space.",
+        default=[],
+    )
+
     args = parser.parse_args()
-    _main(args.root_path)
+    _main(path_to_root=args.root_path, values_to_keep=args.keep)
 
 
 if __name__ == "__main__":

tests/main_test.py

@@ -41,6 +41,18 @@ def is_file(self):
             """
             AWS_PROFILE=""
             AWS_KEY=""
+        """,
+            id="multiple",
+        ),
+        pytest.param(
+            """
+            AWS_PROFILE="123
+            34"
+            AWS_KEY="123"
+        """,
+            """
+            AWS_PROFILE=""
+            AWS_KEY=""
         """,
             id="multiline",
         ),
@@ -84,7 +96,7 @@ def is_file(self):
                 # AWS_SECRET_ACCESS_KEY=""
                 # AWS_SESSION_TOKEN=""
                 """,
-            id="GitHub Issue #1 Example",
+            id="GitHub Issue #3 Example",
         ),
     ),
 )
@@ -102,6 +114,81 @@ def test_clean_function(s, expected):
     assert output == expected
 
 
+@pytest.mark.parametrize(
+    ("s", "expected", "keep"),
+    (
+        pytest.param(
+            "export AWS_PROFILE='test' #exporttest",
+            "export AWS_PROFILE='test' #exporttest",
+            ["AWS_PROFILE"],
+            id="single-keep",
+        ),
+        pytest.param(
+            """
+            AWS_PROFILE="123"
+            AWS_KEY="123"
+        """,
+            """
+            AWS_PROFILE="123"
+            AWS_KEY="123"
+        """,
+            ["AWS_PROFILE", "AWS_KEY"],
+            id="multi-keep",
+        ),
+        pytest.param(
+            """
+            AWS_PROFILE="123
+            34"
+            AWS_KEY="123"
+        """,
+            """
+            AWS_PROFILE="123
+            34"
+            AWS_KEY="123"
+        """,
+            ["AWS_PROFILE", "AWS_KEY"],
+            id="multi-keep",
+        ),
+        pytest.param(
+            """
+            AWS_PROFILE="123"
+            AWS_KEY="123"
+        """,
+            """
+            AWS_PROFILE=""
+            AWS_KEY="123"
+        """,
+            ["aws_profile", "AWS_KEY"],
+            id="case-sensitive",
+        ),
+        pytest.param(
+            """
+            password=pa55w0rd
+            url=www.google.com
+        """,
+            """
+            password=
+            url=www.google.com
+        """,
+            ["url"],
+            id="GitHub Issue #5 Example",
+        ),
+    ),
+)
+def test_clean_function_with_values_to_keep(s, expected, keep):
+    # First we create a temp directory in which we store the .env file
+    tmpdir = tempfile.mkdtemp()
+    # We write the content into a .env
+    with open(f"{tmpdir}/.env", "w") as f:
+        print(s, end="", file=f)
+    clean_dotenv._clean_env(f"{tmpdir}/.env", values_to_keep=keep)
+    # We now get the cleaned file
+    with open(f"{tmpdir}/.env.example", "r") as f:
+        output = f.read()
+    shutil.rmtree(tmpdir)
+    assert output == expected
+
+
 @patch("os.scandir")
 def test_find_dotenv_files(mock_scandir):
     # Mock os.scandir() for files
@@ -131,11 +218,11 @@ def test_find_dotenv_files_function():
 @patch("argparse.ArgumentParser.parse_args")
 @patch("clean_dotenv._main._main")
 def test_main(mock_main, mock_parse_args):
-    mock_parse_args.return_value = MagicMock(root_path="test_rpath")
+    mock_parse_args.return_value = MagicMock(root_path="test_rpath", keep=[])
 
     clean_dotenv.main()
 
-    mock_main.assert_called_once_with("test_rpath")
+    mock_main.assert_called_once_with(path_to_root="test_rpath", values_to_keep=[])
 
 
 def test__main():
@@ -148,11 +235,11 @@ def test__main():
     clean_dotenv._clean_env = mm_clean_env
 
     # Call main method
-    clean_dotenv._main("test_directory")
+    clean_dotenv._main(path_to_root="test_directory", values_to_keep=[])
 
     # Detection should be called once
     mm_find_dotenv.assert_called_once_with("test_directory")
 
     # The creation of new .env file should be called twice, last with "test.env"
     assert mm_clean_env.call_count == 2
-    mm_clean_env.assert_called_with("test.env")
+    mm_clean_env.assert_called_with(path_to_env="test.env", values_to_keep=[])