Repository: theskumar/python-dotenv
Pull request #456 (2023-04-10 16:46:56)

Gracefully handle code which has been imported from a zipfile
I encountered this issue when loading my codebase into Spark via the --py-file option. Basically, if a file in a zip archive tries to call `load_dotenv()`, and that file is then imported, it will cause an error. Hopefully the test cases are self-explanatory!

theskumar: Thanks you @samwyma for this patch and a comprehensive test case. 

Patch:
diff --git a/src/dotenv/main.py b/src/dotenv/main.py
index f40c20e..383b79f 100644
--- a/src/dotenv/main.py
+++ b/src/dotenv/main.py
@@ -291,7 +291,9 @@ def find_dotenv(
         frame = sys._getframe()
         current_file = __file__
 
-        while frame.f_code.co_filename == current_file:
+        while frame.f_code.co_filename == current_file or not os.path.exists(
+            frame.f_code.co_filename
+        ):
             assert frame.f_back is not None
             frame = frame.f_back
         frame_filename = frame.f_code.co_filename


Test patch:
diff --git a/tests/test_zip_imports.py b/tests/test_zip_imports.py
new file mode 100644
index 0000000..46d3c02
--- /dev/null
+++ b/tests/test_zip_imports.py
@@ -0,0 +1,101 @@
+import os
+import sys
+import sh
+import textwrap
+from typing import List
+from unittest import mock
+from zipfile import ZipFile
+
+
+def walk_to_root(path: str):
+    last_dir = None
+    current_dir = path
+    while last_dir != current_dir:
+        yield current_dir
+        (parent_dir, _) = os.path.split(current_dir)
+        last_dir, current_dir = current_dir, parent_dir
+
+
+class FileToAdd:
+    def __init__(self, content: str, path: str):
+        self.content = content
+        self.path = path
+
+
+def setup_zipfile(path, files: List[FileToAdd]):
+    zip_file_path = path / "test.zip"
+    dirs_init_py_added_to = set()
+    with ZipFile(zip_file_path, "w") as zip:
+        for f in files:
+            zip.writestr(data=f.content, zinfo_or_arcname=f.path)
+            for dir in walk_to_root(os.path.dirname(f.path)):
+                if dir not in dirs_init_py_added_to:
+                    print(os.path.join(dir, "__init__.py"))
+                    zip.writestr(
+                        data="", zinfo_or_arcname=os.path.join(dir, "__init__.py")
+                    )
+                    dirs_init_py_added_to.add(dir)
+    return zip_file_path
+
+
+@mock.patch.object(sys, "path", list(sys.path))
+def test_load_dotenv_gracefully_handles_zip_imports_when_no_env_file(tmp_path):
+    zip_file_path = setup_zipfile(
+        tmp_path,
+        [
+            FileToAdd(
+                content=textwrap.dedent(
+                    """
+            from dotenv import load_dotenv
+
+            load_dotenv()
+        """
+                ),
+                path="child1/child2/test.py",
+            ),
+        ],
+    )
+
+    # Should run without an error
+    sys.path.append(str(zip_file_path))
+    import child1.child2.test  # noqa
+
+
+def test_load_dotenv_outside_zip_file_when_called_in_zipfile(tmp_path):
+    zip_file_path = setup_zipfile(
+        tmp_path,
+        [
+            FileToAdd(
+                content=textwrap.dedent(
+                    """
+            from dotenv import load_dotenv
+
+            load_dotenv()
+        """
+                ),
+                path="child1/child2/test.py",
+            ),
+        ],
+    )
+    dotenv_path = tmp_path / ".env"
+    dotenv_path.write_bytes(b"a=b")
+    code_path = tmp_path / "code.py"
+    code_path.write_text(
+        textwrap.dedent(
+            f"""
+        import os
+        import sys
+
+        sys.path.append("{zip_file_path}")
+
+        import child1.child2.test
+
+        print(os.environ['a'])
+    """
+        )
+    )
+    os.chdir(str(tmp_path))
+
+    result = sh.Command(sys.executable)(code_path)
+
+    assert result == "b\n"