Support override variables (theskumar#63)

milonimrod · theskumar · commit 3d84003b251c · 2017-09-08T15:27:02.000+05:30
* Added an override option to the load_dotenv

* Added test to override

* added .idea to .gitignore

* Added readme information

* Added IPython magic option to override

* Added a test for the IPython part. Both with and without override

* Fixed make coverage (now looks for the correct src dir)

* Added the verbose option to ipython

* Added documentation about verbose and override

* Fixed a small typo
diff --git a/.gitignore b/.gitignore
@@ -8,3 +8,4 @@ __pycache__
 .DS_Store
 htmlcov/
 .cache/
+.idea
diff --git a/Makefile b/Makefile
@@ -26,7 +26,7 @@ test:
 	py.test tests/
 
 coverage:
-	coverage run --source=dotenv.py --omit='*tests*' -m py.test tests/ -v --tb=native
+	coverage run --source=dotenv --omit='*tests*' -m py.test tests/ -v --tb=native
 	coverage report
 
 coverage-html: coverage
diff --git a/README.rst b/README.rst
@@ -63,6 +63,12 @@ specified file -- called ``.env`` by default.
     from dotenv import load_dotenv, find_dotenv
     load_dotenv(find_dotenv())
 
+You can also set _load_dotenv_ to override existing variables:
+.. code:: python
+
+    from dotenv import load_dotenv, find_dotenv
+    load_dotenv(find_dotenv(), override=True)
+
 Now, you can access the variables either from system environment
 variable or loaded from ``.env`` file. **System environment variables
 gets higher precedence** and it's advised not to include it in version control.
@@ -154,6 +160,12 @@ You can use dotenv with iPython. You can either let the dotenv search for .env w
     # Specify a particular file
     %dotenv relative/or/absolute/path/to/.env
 
+    # Use _-o_ to indicate override of existing variables
+    %dotenv -o
+
+    # Use _-v_ to turn verbose mode on
+    %dotenv -v
+
 
 Setting config on remote servers
 --------------------------------
@@ -220,7 +232,7 @@ commands like so ``fab config:set,<key1>,<value1> config:set,<key2>,<value2>``
     $ fab config:set,hello,world config:set,foo,bar config:set,fizz=buzz
 
 
-Releated Projects
+Related Projects
 =================
 
 -  `Honcho <https://github.com/nickstenning/honcho>`__ - For managing
diff --git a/dotenv/ipython.py b/dotenv/ipython.py
@@ -1,25 +1,40 @@
 from __future__ import print_function
 from .main import load_dotenv, find_dotenv
 
+from IPython.core.magic import Magics, magics_class, line_magic
+from IPython.core.magic_arguments import (argument, magic_arguments,
+                                          parse_argstring)
 
-def _magic(dotenv_path):
-    """
-    dotenv [dotenv_path]
 
-    Search in increasingly higher folders for the `dotenv_path`
-    """
-    # Locate the .env file
-    dotenv_path = dotenv_path or '.env'
-    try:
-        dotenv_path = find_dotenv(dotenv_path, True, True)
-    except IOError:
-        print("cannot find .env file")
-        return
+@magics_class
+class IPythonDotEnv(Magics):
 
-    # Load the .env file
-    load_dotenv(dotenv_path)
+    @magic_arguments()
+    @argument(
+        '-o', '--override', action='store_true',
+        help="Indicate to override existing variables"
+    )
+    @argument(
+        '-v', '--verbose', action='store_true',
+        help="Indicate function calls to be verbose"
+    )
+    @argument('dotenv_path', nargs='?', type=str, default='.env',
+              help='Search in increasingly higher folders for the `dotenv_path`')
+    @line_magic
+    def dotenv(self, line):
+        args = parse_argstring(self.dotenv, line)
+        # Locate the .env file
+        dotenv_path = args.dotenv_path
+        try:
+            dotenv_path = find_dotenv(dotenv_path, True, True)
+        except IOError:
+            print("cannot find .env file")
+            return
+
+        # Load the .env file
+        load_dotenv(dotenv_path, verbose=args.verbose, override=args.override)
 
 
 def load_ipython_extension(ipython):
     """Register the %dotenv magic."""
-    ipython.register_magic_function(_magic, magic_name='dotenv')
+    ipython.register_magics(IPythonDotEnv)
diff --git a/dotenv/main.py b/dotenv/main.py
@@ -16,7 +16,7 @@ def decode_escaped(escaped):
     return __escape_decoder(escaped)[0]
 
 
-def load_dotenv(dotenv_path, verbose=False):
+def load_dotenv(dotenv_path, verbose=False, override=False):
     """
     Read a .env file and load into os.environ.
     """
@@ -25,7 +25,10 @@ def load_dotenv(dotenv_path, verbose=False):
             warnings.warn("Not loading %s - it doesn't exist." % dotenv_path)
         return None
     for k, v in dotenv_values(dotenv_path).items():
-        os.environ.setdefault(k, v)
+        if override:
+            os.environ[k] = v
+        else:
+            os.environ.setdefault(k, v)
     return True
 
 
diff --git a/requirements.txt b/requirements.txt
@@ -6,3 +6,4 @@ wheel
 pytest-cov
 pytest-flake8
 click
+ipython
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -6,6 +6,7 @@
 import sh
 
 from dotenv import load_dotenv, find_dotenv, set_key
+from IPython.terminal.embed import InteractiveShellEmbed
 
 
 def test_warns_if_file_does_not_exist():
@@ -69,3 +70,43 @@ def test_load_dotenv(cli):
         assert 'DOTENV' in os.environ
         assert os.environ['DOTENV'] == 'WORKS'
         sh.rm(dotenv_path)
+
+
+def test_load_dotenv_override(cli):
+    dotenv_path = '.test_load_dotenv_override'
+    key_name = "DOTENV_OVER"
+
+    with cli.isolated_filesystem():
+        sh.touch(dotenv_path)
+        os.environ[key_name] = "OVERRIDE"
+        set_key(dotenv_path, key_name, 'WORKS')
+        success = load_dotenv(dotenv_path, override=True)
+        assert success
+        assert key_name in os.environ
+        assert os.environ[key_name] == 'WORKS'
+        sh.rm(dotenv_path)
+
+
+def test_ipython():
+    tmpdir = os.path.realpath(tempfile.mkdtemp())
+    os.chdir(tmpdir)
+    filename = os.path.join(tmpdir, '.env')
+    with open(filename, 'w') as f:
+        f.write("MYNEWVALUE=q1w2e3\n")
+    ipshell = InteractiveShellEmbed()
+    ipshell.magic("load_ext dotenv")
+    ipshell.magic("dotenv")
+    assert os.environ["MYNEWVALUE"] == 'q1w2e3'
+
+
+def test_ipython_override():
+    tmpdir = os.path.realpath(tempfile.mkdtemp())
+    os.chdir(tmpdir)
+    filename = os.path.join(tmpdir, '.env')
+    os.environ["MYNEWVALUE"] = "OVERRIDE"
+    with open(filename, 'w') as f:
+        f.write("MYNEWVALUE=q1w2e3\n")
+    ipshell = InteractiveShellEmbed()
+    ipshell.magic("load_ext dotenv")
+    ipshell.magic("dotenv -o")
+    assert os.environ["MYNEWVALUE"] == 'q1w2e3'
