Add method to walk up directories looking for .env (theskumar#23)

isms · theskumar · commit 1fd40c5bcedc · 2016-05-01T19:38:22.000+05:30
- Add `find_dotenv` method that will try to find a .env file by (a) guessing where to start using `__file__` or the working directory -- allowing this to work in non-file contexts such as IPython notebooks and the REPL, and then (b) walking up the dir- ectory tree looking for the specified file -- called `.env` by default. This is a bit like the "filthy magic" employed by django-dotenv[1] to serve the same purpose, and allows the user to write `load_dotenv(find_dotenv())` in many contexts. - Add test for new function [1] https://github.com/jpadilla/django-dotenv/blob/master/dotenv.py#L44-L46
diff --git a/dotenv/__init__.py b/dotenv/__init__.py
@@ -1,4 +1,4 @@
 from .cli import get_cli_string
-from .main import load_dotenv, get_key, set_key, unset_key
+from .main import load_dotenv, get_key, set_key, unset_key, find_dotenv
 
-__all__ = ['get_cli_string', 'load_dotenv', 'get_key', 'set_key', 'unset_key']
+__all__ = ['get_cli_string', 'load_dotenv', 'get_key', 'set_key', 'unset_key', 'find_dotenv']
diff --git a/dotenv/main.py b/dotenv/main.py
@@ -1,5 +1,4 @@
 # -*- coding: utf-8 -*-
-
 import os
 import warnings
 
@@ -95,3 +94,45 @@ def flatten_and_write(dotenv_path, dotenv_as_dict, quote_mode="always"):
             str_format = '%s="%s"\n' if _mode == "always" else '%s=%s\n'
             f.write(str_format % (k, v))
     return True
+
+
+def _walk_to_root(path):
+    """
+    Yield directories starting from the given directory up to the root
+    """
+    if not os.path.exists(path):
+        raise IOError('Starting path not found')
+
+    if os.path.isfile(path):
+        path = os.path.dirname(path)
+
+    last_dir = None
+    current_dir = os.path.abspath(path)
+    while last_dir != current_dir:
+        yield current_dir
+        parent_dir = os.path.abspath(os.path.join(current_dir, os.path.pardir))
+        last_dir, current_dir = current_dir, parent_dir
+
+
+def find_dotenv(filename='.env', raise_error_if_not_found=False, usecwd=False):
+    """
+    Search in increasingly higher folders for the given file
+
+    Returns path to the file if found, or an empty string otherwise
+    """
+    if usecwd or '__file__' not in globals():
+        # should work without __file__, e.g. in REPL or IPython notebook
+        path = os.getcwd()
+    else:
+        # will work for .py files
+        path = os.path.dirname(os.path.abspath(__file__))
+
+    for dirname in _walk_to_root(path):
+        check_path = os.path.join(dirname, filename)
+        if os.path.exists(check_path):
+            return check_path
+
+    if raise_error_if_not_found:
+        raise IOError('File not found')
+
+    return ''
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -1,6 +1,10 @@
+# -*- coding: utf8 -*-
+import os
+import pytest
+import tempfile
 import warnings
 
-from dotenv import load_dotenv
+from dotenv import load_dotenv, find_dotenv
 
 
 def test_warns_if_file_does_not_exist():
@@ -10,3 +14,44 @@ def test_warns_if_file_does_not_exist():
         assert len(w) == 1
         assert w[0].category is UserWarning
         assert str(w[0].message) == "Not loading .does_not_exist - it doesn't exist."
+
+
+def test_find_dotenv():
+    """
+    Create a temporary folder structure like the following:
+
+        tmpXiWxa5/
+        └── child1
+            ├── child2
+            │   └── child3
+            │       └── child4
+            └── .env
+
+    Then try to automatically `find_dotenv` starting in `child4`
+    """
+    tmpdir = tempfile.mkdtemp()
+
+    curr_dir = tmpdir
+    dirs = []
+    for f in ['child1', 'child2', 'child3', 'child4']:
+        curr_dir = os.path.join(curr_dir, f)
+        dirs.append(curr_dir)
+        os.mkdir(curr_dir)
+
+    child1, child4 = dirs[0], dirs[-1]
+
+    # change the working directory for testing
+    os.chdir(child4)
+
+    # try without a .env file and force error
+    with pytest.raises(IOError):
+        find_dotenv(raise_error_if_not_found=True, usecwd=True)
+
+    # try without a .env file and fail silently
+    assert find_dotenv(usecwd=True) == ''
+
+    # now place a .env file a few levels up and make sure it's found
+    filename = os.path.join(child1, '.env')
+    with open(filename, 'w') as f:
+        f.write("TEST=test\n")
+    assert find_dotenv(usecwd=True) == filename
