Refactor move override variable as class variable

theskumar · May 22, 2019 · 7492aba · 7492aba
1 parent 0681a87
commit 7492aba
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 15 deletions.
diff --git a/src/dotenv/main.py b/src/dotenv/main.py
@@ -33,12 +33,13 @@
 
 class DotEnv():
 
-    def __init__(self, dotenv_path, verbose=False, encoding=None):
-        # type: (Union[Text, _PathLike, _StringIO], bool, Union[None, Text]) -> None
+    def __init__(self, dotenv_path, verbose=False, encoding=None, override=False):
+        # type: (Union[Text, _PathLike, _StringIO], bool, Union[None, Text], bool) -> None
         self.dotenv_path = dotenv_path  # type: Union[Text,_PathLike, _StringIO]
         self._dict = None  # type: Optional[Dict[Text, Text]]
         self.verbose = verbose  # type: bool
         self.encoding = encoding  # type: Union[None, Text]
+        self.override = override  # type: bool
 
     @contextmanager
     def _get_stream(self):
@@ -70,18 +71,16 @@ def parse(self):
                 if mapping.key is not None and mapping.value is not None:
                     yield mapping.key, mapping.value
 
-    def set_as_environment_variables(self, override=False):
-        # type: (bool) -> bool
+    def set_as_environment_variables(self):
+        # type: () -> None
         """
         Load the current dotenv as system environemt variable.
         """
         for k, v in self.dict().items():
-            if k in os.environ and not override:
+            if k in os.environ and not self.override:
                 continue
             os.environ[to_env(k)] = to_env(v)
 
-        return True
-
     def get(self, key):
         # type: (Text) -> Optional[Text]
         """
@@ -266,16 +265,17 @@ def find_dotenv(filename='.env', raise_error_if_not_found=False, usecwd=False):
     return ''
 
 
-def load_dotenv(dotenv_path=None, stream=None, verbose=False, override=False, **kwargs):
+def load_dotenv(dotenv_path=None, stream=None, **kwargs):
     # type: (Union[Text, _PathLike, None], Optional[_StringIO], bool, bool, Union[None, Text]) -> bool
     f = dotenv_path or stream or find_dotenv()
-    return DotEnv(f, verbose=verbose, **kwargs).set_as_environment_variables(override=override)
+    dotenv = DotEnv(f, **kwargs).set_as_environment_variables()
+    return dotenv
 
 
-def dotenv_values(dotenv_path=None, stream=None, verbose=False, **kwargs):
+def dotenv_values(dotenv_path=None, stream=None, **kwargs):
     # type: (Union[Text, _PathLike, None], Optional[_StringIO], bool, Union[None, Text]) -> Dict[Text, Text]
     f = dotenv_path or stream or find_dotenv()
-    return DotEnv(f, verbose=verbose, **kwargs).dict()
+    return DotEnv(f, **kwargs).dict()
 
 
 def run_command(command, env):

diff --git a/tests/test_core.py b/tests/test_core.py
@@ -79,8 +79,7 @@ def test_load_dotenv(tmp_path):
     sh.touch(dotenv_path)
     set_key(dotenv_path, 'DOTENV', 'WORKS')
     assert 'DOTENV' not in os.environ
-    success = load_dotenv(dotenv_path)
-    assert success
+    load_dotenv(dotenv_path)
     assert 'DOTENV' in os.environ
     assert os.environ['DOTENV'] == 'WORKS'
 
@@ -92,8 +91,7 @@ def test_load_dotenv_override(tmp_path):
     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
+    load_dotenv(dotenv_path, override=True)
     assert key_name in os.environ
     assert os.environ[key_name] == 'WORKS'