Repository: theskumar/python-dotenv
Pull request #98 (2018-03-03 13:08:56)

Allow `load_dotenv` and `dotenv_values` to work with StringIO objects
supersede #78 

theskumar: @alanjds I've taken your concept and included it in `load_dotenv` and `dotenv_values`. The parse_dotenv was probably not the right place to do it. 

I've also take the liberty to refactor the code.
theskumar: I'll release in v0.8.0
coveralls: 
[![Coverage Status](https://coveralls.io/builds/15790895/badge)](https://coveralls.io/builds/15790895)

Coverage increased (+0.07%) to 86.638% when pulling **4e537a539a8b44f6a2c0dc496e80598eca3e644f on stream-as-input** into **84d16c4edbf36e22fd678f995e367b0f2fc81bb3 on master**.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/15790895/badge)](https://coveralls.io/builds/15790895)

Coverage increased (+0.07%) to 86.638% when pulling **4e537a539a8b44f6a2c0dc496e80598eca3e644f on stream-as-input** into **84d16c4edbf36e22fd678f995e367b0f2fc81bb3 on master**.


Patch:
diff --git a/.gitignore b/.gitignore
index d2a4839..89c02d5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,3 +9,4 @@ __pycache__
 htmlcov/
 .cache/
 .idea
+.pytest_cache/
diff --git a/README.rst b/README.rst
index 52a52fe..33dfc9d 100644
--- a/README.rst
+++ b/README.rst
@@ -28,6 +28,38 @@ production using `12-factor <http://12factor.net/>`__ principles.
 
 Usages
 ======
+The easiest and most common usage consists on calling ``load_dotenv`` when the
+application starts, which will load environment variables from a file named
+``.env`` in the current directory or any of its parents or from the path specificied;
+after that, you can just call the environment-related method you need as
+provided by ``os.getenv``.
+
+``.env`` looks like this:
+
+.. code:: shell
+
+    # a comment and that will be ignored.
+    REDIS_ADDRESS=localhost:6379
+    MEANING_OF_LIFE=42
+    MULTILINE_VAR="hello\nworld"
+
+You can optionally prefix each line with the word ``export``, which will
+conveniently allow you to source the whole file on your shell.
+
+``.env`` can interpolate variables using POSIX variable expansion, variables
+are replaced from the environment first or from other values in the ``.env``
+file if the variable is not present in the environment. (``Note``: Default Value
+Expansion is not supported as of yet, see `#30 <https://github.com/theskumar/python-dotenv/pull/30#issuecomment-244036604>`__.)
+
+.. code:: shell
+
+    CONFIG_PATH=${HOME}/.config/foo
+    DOMAIN=example.org
+    EMAIL=admin@${DOMAIN}
+
+
+Getting started
+================
 
 Assuming you have created the ``.env`` file along-side your settings
 module.
@@ -43,66 +75,64 @@ Add the following code to your ``settings.py``
 .. code:: python
 
     # settings.py
-    from os.path import join, dirname
     from dotenv import load_dotenv
-
-    dotenv_path = join(dirname(__file__), '.env')
-    load_dotenv(dotenv_path)
+    load_dotenv()
 
     # OR, the same with increased verbosity:
-    load_dotenv(dotenv_path, verbose=True)
-
-Alternatively, you can use ``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 directory tree looking for the
-specified file -- called ``.env`` by default.
+    load_dotenv(verbose=True)
 
-.. code:: python
+    # OR, explicitly providing path to '.env'
+    from pathlib import Path  # python3 only
+    env_path = Path('.') / '.env'
+    load_dotenv(dotenv_path=env_path)
 
-    from dotenv import load_dotenv, find_dotenv
-    load_dotenv(find_dotenv())
 
-You can also set ``load_dotenv`` to override existing variables:
+At this point, parsed key/value from the `.env` file is now present as system
+environment variable and they can be conveniently accessed via ``os.getenv()``
 
 .. code:: python
 
-    from dotenv import load_dotenv, find_dotenv
-    load_dotenv(find_dotenv(), override=True)
+    # settings.py
+    import os
+    SECRET_KEY = os.getenv("EMAIL")
+    DATABASE_PASSWORD = os.getenv("DATABASE_PASSWORD")
+
 
-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.
+ ``load_dotenv`` do not override existing System environment variables, but
+ you can do so if you want by passing ``override=True`` to ``load_dotenv()``.
+
+You can use ``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 directory tree looking for the specified file -- called ``.env`` by default.
 
 .. code:: python
 
-    # settings.py
+    from dotenv import load_dotenv, find_dotenv
+    load_dotenv(find_dotenv())
 
-    SECRET_KEY = os.environ.get("SECRET_KEY")
-    DATABASE_PASSWORD = os.environ.get("DATABASE_PASSWORD")
 
+In-memory filelikes
+-------------------
 
-``.env`` is a simple text file. With each environment variables listed
-per line, in the format of ``KEY="Value"``, lines starting with `#` is
-ignored.
+It is possible to not rely on the filesystem to parse filelikes from other sources
+(e.g. from a network storage). ``load_dotenv`` and ``dotenv_values`` accepts a filelike ``stream``.
+Just be sure to rewind it before passing.
 
-.. code:: shell
+.. code:: python
 
-    SOME_VAR=someval
-    # I am a comment and that is OK
-    FOO="BAR"
-    MULTILINE_VAR="hello\nworld"
+    >>> from io import StringIO     # Python2: from StringIO import StringIO
+    >>> from dotenv import dotenv_values
+    >>> filelike = StringIO('SPAM=EGSS\n')
+    >>> filelike.seek(0)
+    >>> parsed = dotenv_values(stream=filelike)
+    >>> parsed['SPAM']
+    'EGSS'
 
-``.env`` can interpolate variables using POSIX variable expansion, variables
-are replaced from the environment first or from other values in the ``.env``
-file if the variable is not present in the environment. (``Note``: Default Value
-Expansion is not supported as of yet, see `#30 <https://github.com/theskumar/python-dotenv/pull/30#issuecomment-244036604>`__.)
 
-.. code:: shell
+The returned value is dictionary with key value pair.
 
-    CONFIG_PATH=${HOME}/.config/foo
-    DOMAIN=example.org
-    EMAIL=admin@${DOMAIN}
+``dotenv_values`` could be useful if you need to *consume* the envfile but not *apply*  it
+directly into the system environment.
 
 
 Django
@@ -112,42 +142,39 @@ If you are using django you should add the above loader script at the
 top of ``wsgi.py`` and ``manage.py``.
 
 
-In-memory filelikes
--------------------
+Installation
+============
 
-Is possible to not rely on the filesystem to parse filelikes from other sources
-(e.g. from a network storage). ``parse_dotenv`` accepts a filelike `stream`.
-Just be sure to rewind it before passing.
+::
 
-.. code:: python
+    pip install -U python-dotenv
 
-    from io import StringIO     # Python2: from StringIO import StringIO
-    from dotenv.main import parse_dotenv
-    filelike = StringIO('SPAM=EGSS\n')
-    filelike.seek(0)
-    parsed = parse_dotenv(stream=filelike)
 
-The returned lazy generator yields ``(key,value)`` tuples.
-To ease the consumption, is suggested to unpack it into a dictionary.
+iPython Support
+---------------
 
-.. code:: python
+You can use dotenv with iPython. You can either let the dotenv search for .env with `%dotenv` or provide the path to .env file explicitly, see below for usages.
 
-    os_env_like = dict(iter(parsed))
-    assert os_env_like['SPAM'] == 'EGGS'
+::
 
-This could be useful if you need to *consume* the envfile but not *apply*  it
-directly into the system environment.
+    %load_ext dotenv
 
+    # Use find_dotenv to locate the file
+    %dotenv
 
-Installation
-============
+    # 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
 
-::
 
-    pip install -U python-dotenv
 
 Command-line interface
-======================
+=================
 
 A cli interface ``dotenv`` is also included, which helps you manipulate
 the ``.env`` file without manually opening it. The same cli installed on
@@ -175,27 +202,6 @@ update your settings on remote server, handy isn't it!
       set    Store the given key/value.
       unset  Removes the given key.
 
-iPython Support
----------------
-
-You can use dotenv with iPython. You can either let the dotenv search for .env with `%dotenv` or provide the path to .env file explicitly, see below for usages.
-
-::
-
-    %load_ext dotenv
-
-    # Use find_dotenv to locate the file
-    %dotenv
-
-    # 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
 --------------------------------
@@ -235,6 +241,7 @@ Get all your remote config info with ``fab config``
 ::
 
     $ fab config
+    foo="bar"
 
 Set remote config variables with ``fab config:set,<key>,<value>``
 
@@ -292,6 +299,13 @@ Executing the tests:
 Changelog
 =========
 
+0.8.0
+----------------------------
+- ``set_key`` and ``unset_key`` only modified the affected file instead of parsing and re-writing file, this causes comments and other file entact as it is.
+- Add support for ``export `` prefix in the line.
+- Internal refractoring (`@theskumar`_)
+- Allow ``load_dotenv`` and ``dotenv_values`` to work with ``StringIO())`` (`@alanjds`_)(`@theskumar`_) (`#78`_)
+
 0.7.1
 ----------------------------
 
@@ -340,6 +354,7 @@ Changelog
 - cli: Added ``-q/--quote`` option to control the behaviour of quotes around values in ``.env``. (Thanks `@hugochinchilla`_).
 - Improved test coverage.
 
+.. _@alanjds: https://github.com/alanjds
 .. _@milonimrod: https://github.com/milonimrod
 .. _@maxkoryukov: https://github.com/maxkoryukov
 .. _@pjona: https://github.com/pjona
@@ -352,6 +367,7 @@ Changelog
 .. _@paulochf: https://github.com/paulochf
 .. _@theskumar: https://github.com/theskumar
 
+.. _#78: https://github.com/theskumar/python-dotenv/issues/78
 .. _#63: https://github.com/theskumar/python-dotenv/issues/63
 .. _#60: https://github.com/theskumar/python-dotenv/issues/60
 .. _#57: https://github.com/theskumar/python-dotenv/issues/57
diff --git a/dotenv/__init__.py b/dotenv/__init__.py
index e766351..60108b7 100644
--- a/dotenv/__init__.py
+++ b/dotenv/__init__.py
@@ -1,8 +1,15 @@
 from .cli import get_cli_string
-from .main import load_dotenv, get_key, set_key, unset_key, find_dotenv
+from .main import load_dotenv, get_key, set_key, unset_key, find_dotenv, dotenv_values
 
 
-__all__ = ['get_cli_string', 'load_dotenv', 'get_key', 'set_key', 'unset_key', 'find_dotenv', 'load_ipython_extension']
+__all__ = ['get_cli_string',
+           'load_dotenv',
+           'dotenv_values',
+           'get_key',
+           'set_key',
+           'unset_key',
+           'find_dotenv',
+           'load_ipython_extension']
 
 
 def load_ipython_extension(ipython):
diff --git a/dotenv/cli.py b/dotenv/cli.py
index bd9bd7a..2c41629 100644
--- a/dotenv/cli.py
+++ b/dotenv/cli.py
@@ -2,7 +2,7 @@ import os
 
 import click
 
-from .main import get_key, dotenv_values, set_key, unset_key
+from .main import dotenv_values, get_key, set_key, unset_key
 
 
 @click.group()
diff --git a/dotenv/compat.py b/dotenv/compat.py
new file mode 100644
index 0000000..c4a481e
--- /dev/null
+++ b/dotenv/compat.py
@@ -0,0 +1,4 @@
+try:
+    from StringIO import StringIO  # noqa
+except ImportError:
+    from io import StringIO  # noqa
diff --git a/dotenv/ipython.py b/dotenv/ipython.py
index 7d84ac2..06252f1 100644
--- a/dotenv/ipython.py
+++ b/dotenv/ipython.py
@@ -1,10 +1,11 @@
 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 import Magics, line_magic, magics_class
 from IPython.core.magic_arguments import (argument, magic_arguments,
                                           parse_argstring)
 
+from .main import find_dotenv, load_dotenv
+
 
 @magics_class
 class IPythonDotEnv(Magics):
diff --git a/dotenv/main.py b/dotenv/main.py
index 9b1efe9..3d1bd72 100644
--- a/dotenv/main.py
+++ b/dotenv/main.py
@@ -1,13 +1,17 @@
 # -*- coding: utf-8 -*-
-from __future__ import absolute_import
+from __future__ import absolute_import, print_function, unicode_literals
 
 import codecs
+import fileinput
+import io
 import os
+import re
 import sys
 import warnings
-import re
 from collections import OrderedDict
 
+from .compat import StringIO
+
 __escape_decoder = codecs.getdecoder('unicode_escape')
 __posix_variable = re.compile('\$\{[^\}]*\}')
 
@@ -16,20 +20,94 @@ def decode_escaped(escaped):
     return __escape_decoder(escaped)[0]
 
 
-def load_dotenv(dotenv_path, verbose=False, override=False):
-    """
-    Read a .env file and load into os.environ.
-    """
-    if not os.path.exists(dotenv_path):
-        if verbose:
-            warnings.warn("Not loading %s - it doesn't exist." % dotenv_path)
-        return None
-    for k, v in dotenv_values(dotenv_path).items():
-        if override:
+def parse_line(line):
+    line = line.strip()
+
+    # Ignore lines with `#` or which doesn't have `=` in it.
+    if not line or line.startswith('#') or '=' not in line:
+        return None, None
+
+    k, v = line.split('=', 1)
+
+    if k.startswith('export '):
+        k = k.lstrip('export ')
+
+    # Remove any leading and trailing spaces in key, value
+    k, v = k.strip(), v.strip()
+
+    if v:
+        v = v.encode('unicode-escape').decode('ascii')
+        quoted = v[0] == v[-1] in ['"', "'"]
+        if quoted:
+            v = decode_escaped(v[1:-1])
+
+    return k, v
+
+
+class DotEnv():
+
+    def __init__(self, dotenv_path, verbose=False):
+        self.dotenv_path = dotenv_path
+        self._dict = None
+        self.verbose = verbose
+
+    def _get_stream(self):
+        self._is_file = False
+        if isinstance(self.dotenv_path, StringIO):
+            return self.dotenv_path
+
+        if os.path.exists(self.dotenv_path):
+            self._is_file = True
+            return io.open(self.dotenv_path)
+
+        if self.verbose:
+            warnings.warn("File doesn't exist {}".format(self.dotenv_path))
+
+        return StringIO('')
+
+    def dict(self):
+        """Return dotenv as dict"""
+        if self._dict:
+            return self._dict
+
+        values = OrderedDict(self.parse())
+        self._dict = resolve_nested_variables(values)
+        return self._dict
+
+    def parse(self):
+        f = self._get_stream()
+
+        for line in f:
+            key, value = parse_line(line)
+            if not key:
+                continue
+
+            yield key, value
+
+        if self._is_file:
+            f.close()
+
+    def set_as_environment_variables(self, override=False):
+        """
+        Load the current dotenv as system environemt variable.
+        """
+        for k, v in self.dict().items():
+            if k in os.environ and not override:
+                continue
             os.environ[k] = v
-        else:
-            os.environ.setdefault(k, v)
-    return True
+
+        return True
+
+    def get(self, key):
+        """
+        """
+        data = self.dict()
+
+        if key in data:
+            return data[key]
+
+        if self.verbose:
+            warnings.warn("key %s not found in %s." % (key, self.dotenv_path))
 
 
 def get_key(dotenv_path, key_to_get):
@@ -38,16 +116,7 @@ def get_key(dotenv_path, key_to_get):
 
     If the .env path given doesn't exist, fails
     """
-    key_to_get = str(key_to_get)
-    if not os.path.exists(dotenv_path):
-        warnings.warn("can't read %s - it doesn't exist." % dotenv_path)
-        return None
-    dotenv_as_dict = dotenv_values(dotenv_path)
-    if key_to_get in dotenv_as_dict:
-        return dotenv_as_dict[key_to_get]
-    else:
-        warnings.warn("key %s not found in %s." % (key_to_get, dotenv_path))
-        return None
+    return DotEnv(dotenv_path, verbose=True).get(key_to_get)
 
 
 def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always"):
@@ -57,15 +126,30 @@ def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always"):
     If the .env path given doesn't exist, fails instead of risking creating
     an orphan .env somewhere in the filesystem
     """
-    key_to_set = str(key_to_set)
-    value_to_set = str(value_to_set).strip("'").strip('"')
+    value_to_set = value_to_set.strip("'").strip('"')
     if not os.path.exists(dotenv_path):
         warnings.warn("can't write to %s - it doesn't exist." % dotenv_path)
         return None, key_to_set, value_to_set
-    dotenv_as_dict = OrderedDict(parse_dotenv(dotenv_path))
-    dotenv_as_dict[key_to_set] = value_to_set
-    success = flatten_and_write(dotenv_path, dotenv_as_dict, quote_mode)
-    return success, key_to_set, value_to_set
+
+    if " " in value_to_set:
+        quote_mode = "always"
+
+    line_template = '{}="{}"' if quote_mode == "always" else '{}={}'
+    line_out = line_template.format(key_to_set, value_to_set)
+
+    replaced = False
+    for line in fileinput.input(dotenv_path, inplace=True):
+        k, v = parse_line(line)
+        if k == key_to_set:
+            replaced = True
+            line = line_out
+        print(line, end='')
+
+    if not replaced:
+        with io.open(dotenv_path, "a") as f:
+            f.write("{}\n".format(line_out))
+
+    return True, key_to_set, value_to_set
 
 
 def unset_key(dotenv_path, key_to_unset, quote_mode="always"):
@@ -75,51 +159,24 @@ def unset_key(dotenv_path, key_to_unset, quote_mode="always"):
     If the .env path given doesn't exist, fails
     If the given key doesn't exist in the .env, fails
     """
-    key_to_unset = str(key_to_unset)
+    removed = False
+
     if not os.path.exists(dotenv_path):
         warnings.warn("can't delete from %s - it doesn't exist." % dotenv_path)
         return None, key_to_unset
-    dotenv_as_dict = dotenv_values(dotenv_path)
-    if key_to_unset in dotenv_as_dict:
-        dotenv_as_dict.pop(key_to_unset, None)
-    else:
-        warnings.warn("key %s not removed from %s - key doesn't exist." % (key_to_unset, dotenv_path))
-        return None, key_to_unset
-    success = flatten_and_write(dotenv_path, dotenv_as_dict, quote_mode)
-    return success, key_to_unset
-
 
-def dotenv_values(dotenv_path):
-    values = OrderedDict(parse_dotenv(dotenv_path))
-    values = resolve_nested_variables(values)
-    return values
-
-
-def parse_dotenv(dotenv_path=None, stream=None):
-    if dotenv_path:
-        f = open(dotenv_path)
-    elif stream:
-        f = stream
-
-    for line in f:
-        line = line.strip()
-        if not line or line.startswith('#') or '=' not in line:
-            continue
-        k, v = line.split('=', 1)
+    for line in fileinput.input(dotenv_path, inplace=True):
+        k, v = parse_line(line)
+        if k == key_to_unset:
+            removed = True
+            line = ''
+        print(line, end='')
 
-        # Remove any leading and trailing spaces in key, value
-        k, v = k.strip(), v.strip().encode('unicode-escape').decode('ascii')
-
-        if len(v) > 0:
-            quoted = v[0] == v[-1] in ['"', "'"]
-
-            if quoted:
-                v = decode_escaped(v[1:-1])
-
-        yield k, v
+    if not removed:
+        warnings.warn("key %s not removed from %s - key doesn't exist." % (key_to_unset, dotenv_path))
+        return None, key_to_unset
 
-    if dotenv_path:
-        f.close()
+    return removed, key_to_unset
 
 
 def resolve_nested_variables(values):
@@ -145,17 +202,6 @@ def resolve_nested_variables(values):
     return values
 
 
-def flatten_and_write(dotenv_path, dotenv_as_dict, quote_mode="always"):
-    with open(dotenv_path, "w") as f:
-        for k, v in dotenv_as_dict.items():
-            _mode = quote_mode
-            if _mode == "auto" and " " in v:
-                _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
@@ -197,3 +243,13 @@ def find_dotenv(filename='.env', raise_error_if_not_found=False, usecwd=False):
         raise IOError('File not found')
 
     return ''
+
+
+def load_dotenv(dotenv_path=None, stream=None, verbose=False, override=False):
+    f = dotenv_path or stream or find_dotenv()
+    return DotEnv(f, verbose=verbose).set_as_environment_variables(override=override)
+
+
+def dotenv_values(dotenv_path=None, stream=None, verbose=False):
+    f = dotenv_path or stream or find_dotenv()
+    return DotEnv(f, verbose=verbose).dict()
diff --git a/setup.cfg b/setup.cfg
index 9e43433..a97aff0 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -9,5 +9,5 @@ exclude = .tox,.git,docs,venv,.venv
 description-file = README.rst
 
 [tool:pytest]
-norecursedirs = .svn _build tmp* dist venv .git .venv
+norecursedirs = .svn _build tmp* dist venv .git .venv venv
 flake8-max-line-length = 120


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index 2b0215c..8ad4f73 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -1,16 +1,14 @@
 # -*- coding: utf8 -*-
+from __future__ import unicode_literals
+
 import os
 import pytest
 import tempfile
 import warnings
 import sh
-try:
-    from StringIO import StringIO
-except ImportError:
-    from io import StringIO
 
-from dotenv import load_dotenv, find_dotenv, set_key
-from dotenv.main import parse_dotenv
+from dotenv import load_dotenv, find_dotenv, set_key, dotenv_values
+from dotenv.compat import StringIO
 from IPython.terminal.embed import InteractiveShellEmbed
 
 
@@ -20,7 +18,7 @@ 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."
+        assert str(w[0].message) == "File doesn't exist .does_not_exist"
 
 
 def test_find_dotenv():
@@ -117,10 +115,17 @@ def test_ipython_override():
     assert os.environ["MYNEWVALUE"] == 'q1w2e3'
 
 
-def test_parse_dotenv_stream():
-    stream = StringIO('DOTENV=WORKS\n')
+def test_dotenv_values_stream():
+    stream = StringIO(u'hello="it works!😃"\nDOTENV=${hello}\n')
     stream.seek(0)
-    parsed_generator = parse_dotenv(stream=stream)
-    parsed_dict = dict(iter(parsed_generator))
+    parsed_dict = dotenv_values(stream=stream)
     assert 'DOTENV' in parsed_dict
-    assert parsed_dict['DOTENV'] == 'WORKS'
+    assert parsed_dict['DOTENV'] == u'it works!😃'
+
+
+def test_dotenv_values_export():
+    stream = StringIO('export foo=bar\n')
+    stream.seek(0)
+    load_dotenv(stream=stream)
+    assert 'foo' in os.environ
+    assert os.environ['foo'] == 'bar'