From f8c8cb24949417f270ae9a7b9f257f4e65c2b9d6 Mon Sep 17 00:00:00 2001 From: Travis Dent Date: Tue, 28 Jan 2025 12:29:39 -0800 Subject: [PATCH] Make .env file parser more robust. Resolves #232. --- agentstack/generation/files.py | 13 +++++++++++-- tests/fixtures/.env | 3 ++- tests/test_generation_files.py | 13 +++++++------ 3 files changed, 20 insertions(+), 9 deletions(-) diff --git a/agentstack/generation/files.py b/agentstack/generation/files.py index e35b446b..db7cbe8d 100644 --- a/agentstack/generation/files.py +++ b/agentstack/generation/files.py @@ -1,4 +1,5 @@ from typing import Optional, Union +import re import string import os, sys import string @@ -39,6 +40,9 @@ class EnvFile: ``` """ + # split the key-value pair on the first '=' character + # allow spaces around the '=' character + RE_PAIR = re.compile(r"^\s*([^\s=]+)\s*=\s*(.*)$") variables: dict[str, str] def __init__(self, filename: str = ENV_FILENAME): @@ -69,8 +73,13 @@ def parse_line(line) -> tuple[str, str]: Pairs are split on the first '=' character, and stripped of whitespace & quotes. Only the last occurrence of a variable is stored. """ - key, value = line.split('=') - return key.strip(), value.strip(string.whitespace + '"') + match = self.RE_PAIR.match(line) + + if not match: + raise ValueError(f"Invalid line in .env file: {line}") + + key, value = match.groups() + return key, value.strip(' "') if os.path.exists(conf.PATH / self._filename): with open(conf.PATH / self._filename, 'r') as f: diff --git a/tests/fixtures/.env b/tests/fixtures/.env index 9197de0d..17b8d8e3 100644 --- a/tests/fixtures/.env +++ b/tests/fixtures/.env @@ -2,4 +2,5 @@ ENV_VAR1=value1 ENV_VAR2=value_ignored ENV_VAR2=value2 -#ENV_VAR3="" \ No newline at end of file +ENV_VAR3 = "12a34b====" +#ENV_VAR4="" \ No newline at end of file diff --git a/tests/test_generation_files.py b/tests/test_generation_files.py index 92f1aa09..8ed151f7 100644 --- a/tests/test_generation_files.py +++ b/tests/test_generation_files.py @@ -89,9 +89,10 @@ def test_read_env(self): shutil.copy(BASE_PATH / "fixtures/.env", self.project_dir / ".env") env = EnvFile() - assert env.variables == {"ENV_VAR1": "value1", "ENV_VAR2": "value2"} + assert env.variables == {"ENV_VAR1": "value1", "ENV_VAR2": "value2", "ENV_VAR3": "12a34b===="} assert env["ENV_VAR1"] == "value1" assert env["ENV_VAR2"] == "value2" + assert env["ENV_VAR3"] == "12a34b====" with self.assertRaises(KeyError) as _: env["ENV_VAR100"] @@ -105,7 +106,7 @@ def test_write_env(self): tmp_data = open(self.project_dir / ".env").read() assert ( tmp_data - == """\nENV_VAR1=value1\nENV_VAR2=value_ignored\nENV_VAR2=value2\n#ENV_VAR3=""\nENV_VAR100=value2""" + == """\nENV_VAR1=value1\nENV_VAR2=value_ignored\nENV_VAR2=value2\nENV_VAR3 = \"12a34b====\"\n#ENV_VAR4=""\nENV_VAR100=value2""" ) def test_write_env_numeric_that_can_be_boolean(self): @@ -116,20 +117,20 @@ def test_write_env_numeric_that_can_be_boolean(self): env.append_if_new("ENV_VAR101", 1) env = EnvFile() # re-read the file - assert env.variables == {"ENV_VAR1": "value1", "ENV_VAR2": "value2", "ENV_VAR100": "0", "ENV_VAR101": "1"} + assert env.variables == {"ENV_VAR1": "value1", "ENV_VAR2": "value2", "ENV_VAR3": "12a34b====", "ENV_VAR100": "0", "ENV_VAR101": "1"} def test_write_env_commented(self): """We should be able to write a commented-out value.""" shutil.copy(BASE_PATH / "fixtures/.env", self.project_dir / ".env") with EnvFile() as env: - env.append_if_new("ENV_VAR3", "value3") + env.append_if_new("ENV_VAR4", "value3") env = EnvFile() # re-read the file - assert env.variables == {"ENV_VAR1": "value1", "ENV_VAR2": "value2", "ENV_VAR3": "value3"} + assert env.variables == {"ENV_VAR1": "value1", "ENV_VAR2": "value2", "ENV_VAR3": "12a34b====", "ENV_VAR4": "value3"} tmp_file = open(self.project_dir / ".env").read() assert ( tmp_file - == """\nENV_VAR1=value1\nENV_VAR2=value_ignored\nENV_VAR2=value2\n#ENV_VAR3=""\nENV_VAR3=value3""" + == """\nENV_VAR1=value1\nENV_VAR2=value_ignored\nENV_VAR2=value2\nENV_VAR3 = \"12a34b====\"\n#ENV_VAR4=""\nENV_VAR4=value3""" )