From e0e2fadc96bd09b05605d5b4a54c4f2871eb55d7 Mon Sep 17 00:00:00 2001
From: Eli Schwartz <eschwartz93@gmail.com>
Date: Thu, 23 Nov 2023 13:57:51 -0500
Subject: [PATCH] wip: make files() be enter-able

This restores compatibility more or less with the old path() approach
that supports directly running `with x as f: do_stuff(f)`. There is no
more need to use as_file() by hand every time you need to have a file.

Namespace packages are TODO. But currently they are half-broken anyway,
in the sense that namespace packages do not support being inside zip
files unless you use pkg_resources style and importlib_resources does
not support this anyway.

We specifically go to efforts to ensure that the return value of files()
is still usable as-is, representing an isinstance()-compatible
pathlib.Path or zipfile.Path. We do this by replacing it at the last
second with a subclass that adds the context manager magic.
---
 importlib_resources/_common.py | 12 ++++++++++--
 importlib_resources/readers.py | 17 +++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/importlib_resources/_common.py b/importlib_resources/_common.py
index 3c6de1cf..c5cd43be 100644
--- a/importlib_resources/_common.py
+++ b/importlib_resources/_common.py
@@ -12,7 +12,8 @@
 from typing import Union, Optional, cast
 from .abc import ResourceReader, Traversable
 
-from ._compat import wrap_spec
+from ._compat import ZipPath, wrap_spec
+from .readers import EnterablePath, EnterableZip
 
 Package = Union[types.ModuleType, str]
 Anchor = Package
@@ -53,7 +54,14 @@ def files(anchor: Optional[Anchor] = None) -> Traversable:
     """
     Get a Traversable resource for an anchor.
     """
-    return from_package(resolve(anchor))
+    f = from_package(resolve(anchor))
+    if isinstance(f, ZipPath):
+        return EnterableZip(f.root, f.at)
+    elif isinstance(f, pathlib.Path):
+        return EnterablePath(f)
+    else:
+        # `with f as g:` will fail for MultiplexedPath
+        return f
 
 
 def get_resource_reader(package: types.ModuleType) -> Optional[ResourceReader]:
diff --git a/importlib_resources/readers.py b/importlib_resources/readers.py
index 1e2d1bac..ce12eb6c 100644
--- a/importlib_resources/readers.py
+++ b/importlib_resources/readers.py
@@ -170,3 +170,20 @@ def resource_path(self, resource):
 
     def files(self):
         return self.path
+
+
+class EnterablePath(pathlib.Path):
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+class EnterableZip(ZipPath):
+    def __enter__(self):
+        from ._common import as_file
+        self._as_file = as_file(self)
+        return self._as_file.__enter__()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._as_file.__exit__(exc_type, exc_value, traceback)