Allow any text stream (IO[str]) as stream

CHANGELOG.md

@@ -7,16 +7,19 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
-### Added
-
-- The `dotenv_path` argument of `set_key` and `unset_key` now has a type of `Union[str,
-  os.PathLike]` instead of just `os.PathLike` (#347 by [@bbc2]).
-
 ### Changed
 
 - Require Python 3.5 or a later version.  Python 2 and 3.4 are no longer supported. (#341
   by [@bbc2]).
 
+### Added
+
+- The `dotenv_path` argument of `set_key` and `unset_key` now has a type of `Union[str,
+  os.PathLike]` instead of just `os.PathLike` (#347 by [@bbc2]).
+- The `stream` argument of `load_dotenv` and `dotenv_values` can now be a text stream
+  (`IO[str]`), which includes values like `io.StringIO("foo")` and `open("file.env",
+  "r")` (#348 by [@bbc2]).
+
 ## [0.18.0] - 2021-06-20
 
 ### Changed

src/dotenv/main.py

@@ -33,13 +33,15 @@ def with_warn_for_invalid_lines(mappings: Iterator[Binding]) -> Iterator[Binding
 class DotEnv():
     def __init__(
         self,
-        dotenv_path: Union[str, _PathLike, io.StringIO],
+        dotenv_path: Optional[Union[str, _PathLike]],
+        stream: Optional[IO[str]] = None,
         verbose: bool = False,
         encoding: Union[None, str] = None,
         interpolate: bool = True,
         override: bool = True,
     ) -> None:
-        self.dotenv_path = dotenv_path  # type: Union[str,_PathLike, io.StringIO]
+        self.dotenv_path = dotenv_path  # type: Optional[Union[str, _PathLike]]
+        self.stream = stream  # type: Optional[IO[str]]
         self._dict = None  # type: Optional[Dict[str, Optional[str]]]
         self.verbose = verbose  # type: bool
         self.encoding = encoding  # type: Union[None, str]
@@ -48,14 +50,17 @@ def __init__(
 
     @contextmanager
     def _get_stream(self) -> Iterator[IO[str]]:
-        if isinstance(self.dotenv_path, io.StringIO):
-            yield self.dotenv_path
-        elif os.path.isfile(self.dotenv_path):
+        if self.dotenv_path and os.path.isfile(self.dotenv_path):
             with io.open(self.dotenv_path, encoding=self.encoding) as stream:
                 yield stream
+        elif self.stream is not None:
+            yield self.stream
         else:
             if self.verbose:
-                logger.info("Python-dotenv could not find configuration file %s.", self.dotenv_path or '.env')
+                logger.info(
+                    "Python-dotenv could not find configuration file %s.",
+                    self.dotenv_path or '.env',
+                )
             yield io.StringIO('')
 
     def dict(self) -> Dict[str, Optional[str]]:
@@ -290,7 +295,7 @@ def _is_interactive():
 
 def load_dotenv(
     dotenv_path: Union[str, _PathLike, None] = None,
-    stream: Optional[io.StringIO] = None,
+    stream: Optional[IO[str]] = None,
     verbose: bool = False,
     override: bool = False,
     interpolate: bool = True,
@@ -299,7 +304,8 @@ def load_dotenv(
     """Parse a .env file and then load all the variables found as environment variables.
 
     - *dotenv_path*: absolute or relative path to .env file.
-    - *stream*: `StringIO` object with .env content, used if `dotenv_path` is `None`.
+    - *stream*: Text stream (such as `io.StringIO`) with .env content, used if
+      `dotenv_path` is `None`.
     - *verbose*: whether to output a warning the .env file is missing. Defaults to
       `False`.
     - *override*: whether to override the system environment variables with the variables
@@ -308,9 +314,12 @@ def load_dotenv(
 
     If both `dotenv_path` and `stream`, `find_dotenv()` is used to find the .env file.
     """
-    f = dotenv_path or stream or find_dotenv()
+    if dotenv_path is None and stream is None:
+        dotenv_path = find_dotenv()
+
     dotenv = DotEnv(
-        f,
+        dotenv_path=dotenv_path,
+        stream=stream,
         verbose=verbose,
         interpolate=interpolate,
         override=override,
@@ -321,7 +330,7 @@ def load_dotenv(
 
 def dotenv_values(
     dotenv_path: Union[str, _PathLike, None] = None,
-    stream: Optional[io.StringIO] = None,
+    stream: Optional[IO[str]] = None,
     verbose: bool = False,
     interpolate: bool = True,
     encoding: Optional[str] = "utf-8",
@@ -338,9 +347,12 @@ def dotenv_values(
 
     If both `dotenv_path` and `stream`, `find_dotenv()` is used to find the .env file.
     """
-    f = dotenv_path or stream or find_dotenv()
+    if dotenv_path is None and stream is None:
+        dotenv_path = find_dotenv()
+
     return DotEnv(
-        f,
+        dotenv_path=dotenv_path,
+        stream=stream,
         verbose=verbose,
         interpolate=interpolate,
         override=True,

tests/test_main.py

@@ -277,7 +277,7 @@ def test_load_dotenv_redefine_var_used_in_file_with_override(dotenv_file):
 
 
 @mock.patch.dict(os.environ, {}, clear=True)
-def test_load_dotenv_utf_8():
+def test_load_dotenv_string_io_utf_8():
     stream = io.StringIO("a=à")
 
     result = dotenv.load_dotenv(stream=stream)
@@ -286,6 +286,18 @@ def test_load_dotenv_utf_8():
     assert os.environ == {"a": "à"}
 
 
+@mock.patch.dict(os.environ, {}, clear=True)
+def test_load_dotenv_file_stream(dotenv_file):
+    with open(dotenv_file, "w") as f:
+        f.write("a=b")
+
+    with open(dotenv_file, "r") as f:
+        result = dotenv.load_dotenv(stream=f)
+
+    assert result is True
+    assert os.environ == {"a": "b"}
+
+
 def test_load_dotenv_in_current_dir(tmp_path):
     dotenv_path = tmp_path / '.env'
     dotenv_path.write_bytes(b'a=b')
@@ -353,11 +365,21 @@ def test_dotenv_values_file(dotenv_file):
         ({}, "a=b\nc=${a}\nd=e\nc=${d}", True, {"a": "b", "c": "e", "d": "e"}),
     ],
 )
-def test_dotenv_values_stream(env, string, interpolate, expected):
+def test_dotenv_values_string_io(env, string, interpolate, expected):
     with mock.patch.dict(os.environ, env, clear=True):
         stream = io.StringIO(string)
         stream.seek(0)
 
         result = dotenv.dotenv_values(stream=stream, interpolate=interpolate)
 
         assert result == expected
+
+
+def test_dotenv_values_file_stream(dotenv_file):
+    with open(dotenv_file, "w") as f:
+        f.write("a=b")
+
+    with open(dotenv_file, "r") as f:
+        result = dotenv.dotenv_values(stream=f)
+
+    assert result == {"a": "b"}