diff --git a/clean_dotenv/_main.py b/clean_dotenv/_main.py index ef0b294..08f3acd 100644 --- a/clean_dotenv/_main.py +++ b/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__": diff --git a/tests/main_test.py b/tests/main_test.py index 4e777f6..b132ac4 100644 --- a/tests/main_test.py +++ b/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=[])