Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Turn get_stream into a context manager
This avoids the use of the `is_file` class variable by abstracting away
the difference between `StringIO` and a file stream.
  • Loading branch information
bbc2 committed Nov 14, 2018
commit 950a0e7bc63526052bce4eb69842ee6bd7ae3bfe
37 changes: 16 additions & 21 deletions dotenv/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from subprocess import Popen
import warnings
from collections import OrderedDict
from contextlib import contextmanager

from .compat import StringIO, PY2, WIN, text_type

Expand Down Expand Up @@ -52,19 +53,17 @@ def __init__(self, dotenv_path, verbose=False):
self._dict = None
self.verbose = verbose

@contextmanager
def _get_stream(self):
self._is_file = False
if isinstance(self.dotenv_path, StringIO):
return self.dotenv_path

if os.path.isfile(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('')
yield self.dotenv_path
elif os.path.isfile(self.dotenv_path):
with io.open(self.dotenv_path) as stream:
yield stream
else:
if self.verbose:
warnings.warn("File doesn't exist {}".format(self.dotenv_path))
yield StringIO('')

def dict(self):
"""Return dotenv as dict"""
Expand All @@ -76,17 +75,13 @@ def dict(self):
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
with self._get_stream() as stream:
for line in stream:
key, value = parse_line(line)
if not key:
continue

if self._is_file:
f.close()
yield key, value

def set_as_environment_variables(self, override=False):
"""
Expand Down