From ae5b3eeefbc63e39ff331af72587a69fc2925ea0 Mon Sep 17 00:00:00 2001
From: Jake Howard <git@theorangeone.net>
Date: Tue, 18 Dec 2018 21:42:23 +0000
Subject: [PATCH] Add env file parsing

---
 catfish/project/__init__.py              | 13 ++++++++++++-
 catfish/worker/server.py                 |  1 +
 example/etc/environments/development/env |  1 +
 setup.py                                 |  1 +
 tests/test_project/test_project.py       |  3 +++
 tests/test_worker/test_server.py         |  8 ++++++++
 6 files changed, 26 insertions(+), 1 deletion(-)
 create mode 100644 example/etc/environments/development/env

diff --git a/catfish/project/__init__.py b/catfish/project/__init__.py
index 46ff790..aff42b5 100644
--- a/catfish/project/__init__.py
+++ b/catfish/project/__init__.py
@@ -1,9 +1,12 @@
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import List
+from typing import Dict, List
+
+from dotenv import dotenv_values
 
 PROCFILE_LOCATIONS = ["Procfile", "etc/environments/development/Procfile"]
+ENVFILE_LOCATIONS = [".env", "etc/environments/development/env"]
 PROCFILE_LINE = re.compile(r"^([A-Za-z0-9_]+):\s*(.+)$")
 
 
@@ -26,11 +29,13 @@ def parse_procfile_processes(project, procfile_lines):
 class Project:
     root: Path
     processes: List["Process"] = field(default_factory=list)
+    env: Dict[str, str] = field(default_factory=dict)
 
     def __post_init__(self):
         self.root = Path(self.root)
         assert self.exists()
         self.processes = self.read_processes()
+        self.read_envfiles()
 
     def read_processes(self):
         for location in PROCFILE_LOCATIONS:
@@ -40,6 +45,12 @@ class Project:
                     return list(parse_procfile_processes(self, f.readlines()))
         return []
 
+    def read_envfiles(self):
+        for location in ENVFILE_LOCATIONS:
+            envfile_path = self.root.joinpath(location)
+            if envfile_path.exists():
+                self.env.update(dotenv_values(envfile_path))
+
     def exists(self):
         return self.root.exists() and self.root.is_dir()
 
diff --git a/catfish/worker/server.py b/catfish/worker/server.py
index d0c3ef7..7400b41 100644
--- a/catfish/worker/server.py
+++ b/catfish/worker/server.py
@@ -70,6 +70,7 @@ async def run_process_command(project: Project, process: Process):
         stderr=subprocess.STDOUT,
         env={
             **os.environ,
+            **project.env,
             "PYTHONUNBUFFERED": "1",
             "CATFISH_IDENT": process.ident,
             "PATH": "{}:{}".format(
diff --git a/example/etc/environments/development/env b/example/etc/environments/development/env
new file mode 100644
index 0000000..c075a74
--- /dev/null
+++ b/example/etc/environments/development/env
@@ -0,0 +1 @@
+FOO=bar
diff --git a/setup.py b/setup.py
index cc6f0f3..febb2e5 100644
--- a/setup.py
+++ b/setup.py
@@ -19,6 +19,7 @@ setup(
         "ujson",
         "pyzmq",
         "aiofiles",
+        "python-dotenv",
     ],
     entry_points="""
         [console_scripts]
diff --git a/tests/test_project/test_project.py b/tests/test_project/test_project.py
index 714befa..1d1a11e 100644
--- a/tests/test_project/test_project.py
+++ b/tests/test_project/test_project.py
@@ -34,6 +34,9 @@ class ProjectTestCase(BaseTestCase):
     def test_name(self):
         self.assertEqual(self.project.name, "example")
 
+    def test_read_environment(self):
+        self.assertEqual(self.project.env["FOO"], "bar")
+
 
 class ProcessTestCase(BaseTestCase):
     def setUp(self):
diff --git a/tests/test_worker/test_server.py b/tests/test_worker/test_server.py
index 91fe445..5ae84e8 100644
--- a/tests/test_worker/test_server.py
+++ b/tests/test_worker/test_server.py
@@ -34,6 +34,14 @@ class ProcessWorkerTestCase(BaseWorkerTestCase):
         env = psutil.Process(response["pid"]).environ()
         self.assertEqual(env["PYTHONUNBUFFERED"], "1")
         self.assertEqual(env["CATFISH_IDENT"], self.process.ident)
+        self.assertEqual(env["FOO"], "bar")
+
+    def test_additional_path(self):
+        response = send_to_server(
+            PayloadType.PROCESS,
+            {"path": str(self.project.root), "process": str(self.process.name)},
+        )
+        env = psutil.Process(response["pid"]).environ()
         path_dirs = env["PATH"].split(":")
         for path in self.project.get_extra_path():
             self.assertIn(str(path), path_dirs)