Rework our machinery for `importlib.resources`

commit_hash:952ba013b771d9c6cb949cf43125956ad5cdfd58

5 files changed, 141 insertions, 113 deletions

diff --git a/library/python/runtime_py3/__res.py b/library/python/runtime_py3/__res.py
index c2f94a4b245..efbef0d564a 100644
--- a/library/python/runtime_py3/__res.py
+++ b/library/python/runtime_py3/__res.py
@@ -527,9 +527,10 @@ class ResourceImporter(SourceFileLoader):
                 yield m
 
     def get_resource_reader(self, fullname):
-        import os
-        path = os.path.dirname(self.get_filename(fullname))
-        return _ResfsResourceReader(self, path)
+        """Return the ResourceReader for a module in a binary file."""
+        from sitecustomize import ResfsTraversableResources
+
+        return ResfsTraversableResources(self, fullname)
 
     @staticmethod
     def find_distributions(*args, **kwargs):
@@ -542,52 +543,8 @@ class ResourceImporter(SourceFileLoader):
         of directories ``context.path``.
         """
         from sitecustomize import MetadataArcadiaFinder
-        return MetadataArcadiaFinder.find_distributions(*args, **kwargs)
-
-
-class _ResfsResourceReader:
-
-    def __init__(self, importer, path):
-        self.importer = importer
-        self.path = path
-
-    def open_resource(self, resource):
-        path = f'{self.path}/{resource}'
-        from io import BytesIO
-        try:
-            return BytesIO(self.importer.get_data(path))
-        except OSError:
-            raise FileNotFoundError(path)
 
-    def resource_path(self, resource):
-        # All resources are in the binary file, so there is no path to the file.
-        # Raising FileNotFoundError tells the higher level API to extract the
-        # binary data and create a temporary file.
-        raise FileNotFoundError
-
-    def is_resource(self, name):
-        path = f'{self.path}/{name}'
-        try:
-            self.importer.get_data(path)
-        except OSError:
-            return False
-        return True
-
-    def contents(self):
-        subdirs_seen = set()
-        len_path = len(self.path) + 1  # path + /
-        for key in resfs_files(f"{self.path}/"):
-            relative = key[len_path:]
-            res_or_subdir, *other = relative.split(b'/')
-            if not other:
-                yield _s(res_or_subdir)
-            elif res_or_subdir not in subdirs_seen:
-                subdirs_seen.add(res_or_subdir)
-                yield _s(res_or_subdir)
-
-    def files(self):
-        import sitecustomize
-        return sitecustomize.ArcadiaResourceContainer(f"resfs/file/{self.path}/")
+        return MetadataArcadiaFinder.find_distributions(*args, **kwargs)
 
 
 class ArcadiaSourceFinder:
diff --git a/library/python/runtime_py3/sitecustomize.py b/library/python/runtime_py3/sitecustomize.py
index dbd134f5a47..db6f0ca6f05 100644
--- a/library/python/runtime_py3/sitecustomize.py
+++ b/library/python/runtime_py3/sitecustomize.py
@@ -7,16 +7,20 @@ import stat
 import sys
 import typing
 
+from collections.abc import Iterator
+from importlib.abc import Loader
 from importlib.metadata import (
     Distribution,
     DistributionFinder,
     Prepared,
 )
-from importlib.resources.abc import Traversable
+from importlib.resources.abc import Traversable, TraversableResources
 
 from __res import find, iter_keys, resfs_files, resfs_read
 
-METADATA_NAME = re.compile("^Name: (.*)$", re.MULTILINE)
+METADATA_NAME: typing.Final[re.Pattern[str]] = re.compile("^Name: (.*)$", re.MULTILINE)
+
+RESFS_PREFIX: typing.Final[typing.Literal["resfs/file/"]] = "resfs/file/"
 
 try:
     BINARY_STAT: typing.Final[os.stat_result] = os.stat(sys.executable)
@@ -62,62 +66,112 @@ RESOURCE_DIRECTORY_STAT: typing.Final[os.stat_result] = os.stat_result(
 )
 
 
-class ArcadiaTraversable(Traversable):
-    def __init__(self, resfs) -> None:
-        self._resfs = str(resfs)
-        self._path = pathlib.Path(resfs)
+class ResfsTraversableResources(TraversableResources):
+
+    __slots__ = ("_resfs_prefix",)
+
+    def __init__(self, loader: Loader, fullname: str) -> None:
+        filename = loader.get_filename(fullname).replace("\\", "/")
+        self._resfs_prefix = f"{RESFS_PREFIX}{os.path.dirname(filename)}/"
+
+    def _get_data(self, resource: str) -> bytes:
+        path = f"{self._resfs_prefix}{resource}"
+        data = find(path.encode("utf-8"))
+        if data is not None:
+            return data
+        raise FileNotFoundError(path)
+
+    def open_resource(self, resource: str) -> io.BytesIO:
+        return io.BytesIO(self._get_data(resource))
+
+    def is_resource(self, path: str) -> bool:
+        try:
+            self._get_data(path)
+        except FileNotFoundError:
+            return False
+        else:
+            return True
+
+    def contents(self) -> Iterator[str]:
+        subdirs_seen = set()
+        len_prefix = len(self._resfs_prefix) - len(RESFS_PREFIX)
+        for key in resfs_files(self._resfs_prefix.removeprefix(RESFS_PREFIX)):
+            res_or_subdir, *other = key[len_prefix:].split(b"/")
+            if not other:
+                yield res_or_subdir.decode("utf-8")
+            elif res_or_subdir not in subdirs_seen:
+                subdirs_seen.add(res_or_subdir)
+                yield res_or_subdir.decode("utf-8")
+
+    def files(self) -> Traversable:
+        return ResfsResourceContainer(self._resfs_prefix)
+
+
+class ResfsTraversable(Traversable):
+
+    __slots__ = ("_resfs",)
+
+    def __init__(self, resfs_path: str | pathlib.Path) -> None:
+        self._resfs = str(resfs_path)
+
+    @functools.cached_property
+    def _resfs_path(self) -> pathlib.Path:
+        return pathlib.Path(self._resfs)
 
     def __eq__(self, other) -> bool:
-        if isinstance(other, ArcadiaTraversable):
-            return self._path == other._path
+        if isinstance(other, ResfsTraversable):
+            return self._resfs_path == other._resfs_path
         raise NotImplementedError
 
     def __lt__(self, other) -> bool:
-        if isinstance(other, ArcadiaTraversable):
-            return self._path < other._path
+        if isinstance(other, ResfsTraversable):
+            return self._resfs_path < other._resfs_path
         raise NotImplementedError
 
     def __hash__(self) -> int:
-        return hash(self._path)
+        return hash(self._resfs_path)
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self._resfs!r})"
 
     @property
     def name(self) -> str:
-        return self._path.name
+        return self._resfs_path.name
 
     @property
     def suffix(self) -> str:
-        return self._path.suffix
+        return self._resfs_path.suffix
 
     @property
     def stem(self) -> str:
-        return self._path.stem
+        return self._resfs_path.stem
 
     def with_suffix(self, suffix: str) -> typing.Self:
-        return type(self)(self._path.with_suffix(suffix))
+        return type(self)(self._resfs_path.with_suffix(suffix))
 
     def lstat(self) -> os.stat_result:
         return self.stat()
 
-    def relative_to(self, other):
-        if isinstance(other, ArcadiaTraversable):
-            return self._path.relative_to(other._path)
+    def relative_to(self, other) -> pathlib.Path:
+        if isinstance(other, ResfsTraversable):
+            return self._resfs_path.relative_to(other._resfs_path)
         raise NotImplementedError
 
     def resolve(self, strict: bool = False) -> typing.Self:
         return self
 
 
-class ArcadiaResource(ArcadiaTraversable):
+class ResfsResource(ResfsTraversable):
     def is_file(self) -> bool:
         return True
 
     def is_dir(self) -> bool:
         return False
 
-    def open(self, mode="r", *args, **kwargs):
+    def iterdir(self) -> Iterator[ResfsTraversable]:
+        return iter(())
+
+    def open(self, mode="r", *args, **kwargs) -> typing.IO:
         data = self.data
         if data is None:
             raise FileNotFoundError(self._resfs)
@@ -129,12 +183,6 @@ class ArcadiaResource(ArcadiaTraversable):
 
         return stream
 
-    def joinpath(self, *name) -> ArcadiaTraversable:
-        raise RuntimeError("Cannot traverse into a resource")
-
-    def iterdir(self):
-        return iter(())
-
     @functools.cached_property
     def data(self) -> bytes | None:
         return find(self._resfs.encode("utf-8"))
@@ -165,50 +213,28 @@ class ArcadiaResource(ArcadiaTraversable):
         )
 
 
-class ArcadiaResourceContainer(ArcadiaTraversable):
-    def is_dir(self) -> bool:
-        return True
-
+class ResfsResourceContainer(ResfsTraversable):
     def is_file(self) -> bool:
         return False
 
-    def iterdir(self):
+    def is_dir(self) -> bool:
+        return True
+
+    def iterdir(self) -> Iterator[ResfsTraversable]:
         seen = set()
         for key, path_without_prefix in iter_keys(self._resfs.encode("utf-8")):
             if b"/" in path_without_prefix:
-                subdir = path_without_prefix.split(b"/", maxsplit=1)[0].decode("utf-8")
+                subdir = path_without_prefix.split(b"/", maxsplit=1)[0]
                 if subdir not in seen:
                     seen.add(subdir)
-                    yield ArcadiaResourceContainer(f"{self._resfs}{subdir}/")
+                    yield ResfsResourceContainer(f"{self._resfs}{subdir.decode("utf-8")}/")
             else:
-                yield ArcadiaResource(key.decode("utf-8"))
+                yield ResfsResource(key.decode("utf-8"))
 
-    def open(self, *args, **kwargs):
+    def open(self, *args, **kwargs) -> typing.IO:
         raise IsADirectoryError(self._resfs)
 
     @staticmethod
-    def _flatten(compound_names):
-        for name in compound_names:
-            yield from str(name).split("/")
-
-    def joinpath(self, *descendants) -> ArcadiaTraversable:
-        if not descendants:
-            return self
-
-        names = self._flatten(descendants)
-        target = next(names)
-        if target == ".":
-            return self
-        for traversable in self.iterdir():
-            if traversable.name == target:
-                if isinstance(traversable, ArcadiaResource):
-                    return traversable
-                else:
-                    return traversable.joinpath(*names)
-
-        raise FileNotFoundError("/".join(self._flatten(descendants)))
-
-    @staticmethod
     def stat() -> os.stat_result:
         return RESOURCE_DIRECTORY_STAT
 
diff --git a/library/python/runtime_py3/test/resources/data/my_data b/library/python/runtime_py3/test/resources/data/my_data
new file mode 100644
index 00000000000..6320cd248dd
--- /dev/null
+++ b/library/python/runtime_py3/test/resources/data/my_data
@@ -0,0 +1 @@
+data
\ No newline at end of file
diff --git a/library/python/runtime_py3/test/test_resources.py b/library/python/runtime_py3/test/test_resources.py
index 0bf70f8c54b..fda82527886 100644
--- a/library/python/runtime_py3/test/test_resources.py
+++ b/library/python/runtime_py3/test/test_resources.py
@@ -1,7 +1,12 @@
 import importlib.resources as ir
 
+from importlib.resources._common import get_resource_reader
+from importlib.resources.abc import TraversalError
+
 import pytest
 
+from sitecustomize import ResfsTraversableResources
+
 
 @pytest.mark.parametrize(
     "package, resource",
@@ -25,9 +30,16 @@ def test_is_resource_missing(package, resource):
     assert not ir.is_resource(package, resource)
 
 
-def test_is_resource_subresource_directory():
[email protected](
+    "directory",
+    (
+        "data",
+        "submodule",
+    ),
+)
+def test_is_resource_subresource_directory(directory):
     # Directories are not resources.
-    assert not ir.is_resource("resources", "submodule")
+    assert not ir.is_resource("resources", directory)
 
 
 @pytest.mark.parametrize(
@@ -42,7 +54,7 @@ def test_read_binary_good_path(package, resource, expected):
 
 
 def test_read_binary_missing():
-    with pytest.raises(FileNotFoundError):
+    with pytest.raises(TraversalError):
         ir.read_binary("resources", "111.txt")
 
 
@@ -58,14 +70,14 @@ def test_read_text_good_path(package, resource, expected):
 
 
 def test_read_text_missing():
-    with pytest.raises(FileNotFoundError):
+    with pytest.raises(TraversalError):
         ir.read_text("resources", "111.txt")
 
 
 @pytest.mark.parametrize(
     "package, expected",
     (
-        ("resources", ["submodule", "foo.txt"]),
+        ("resources", ["data", "submodule", "foo.txt"]),
         ("resources.submodule", ["bar.txt"]),
     ),
 )
@@ -76,6 +88,7 @@ def test_contents_good_path(package, expected):
 def test_files_joinpath():
     assert ir.files("resources") / "submodule"
     assert ir.files("resources") / "foo.txt"
+    assert ir.files("resources") / "data" / "my_data"
     assert ir.files("resources") / "submodule" / "bar.txt"
     assert ir.files("resources.submodule") / "bar.txt"
 
@@ -84,6 +97,8 @@ def test_files_joinpath():
     "package, resource, expected",
     (
         ("resources", "foo.txt", b"bar"),
+        ("resources", "data/my_data", b"data"),
+        ("resources", "submodule/bar.txt", b"foo"),
         ("resources.submodule", "bar.txt", b"foo"),
     ),
 )
@@ -95,6 +110,8 @@ def test_files_read_bytes(package, resource, expected):
     "package, resource, expected",
     (
         ("resources", "foo.txt", "bar"),
+        ("resources", "data/my_data", "data"),
+        ("resources", "submodule/bar.txt", "foo"),
         ("resources.submodule", "bar.txt", "foo"),
     ),
 )
@@ -105,7 +122,7 @@ def test_files_read_text(package, resource, expected):
 @pytest.mark.parametrize(
     "package, expected",
     (
-        ("resources", ("foo.txt", "submodule")),
+        ("resources", ("foo.txt", "data", "submodule")),
         ("resources.submodule", ("bar.txt",)),
     ),
 )
@@ -116,9 +133,33 @@ def test_files_iterdir(package, expected):
 @pytest.mark.parametrize(
     "package, expected",
     (
-        ("resources", ("foo.txt", "submodule")),
+        ("resources", ("data", "foo.txt", "submodule")),
         ("resources.submodule", ("bar.txt",)),
     ),
 )
 def test_files_iterdir_with_sort(package, expected):
     assert tuple(resource.name for resource in sorted(ir.files(package).iterdir())) == expected
+
+
+def test_get_resource_reader():
+    import resources
+
+    reader = get_resource_reader(resources)
+    assert isinstance(reader, ResfsTraversableResources)
+
+    assert reader.is_resource("foo.txt") is True
+    assert reader.is_resource("submodule/bar.txt") is True
+    assert reader.is_resource("notfound.txt") is False
+    assert reader.is_resource("submodule") is False
+
+    with pytest.raises(FileNotFoundError) as ex:
+        reader.resource_path("foo.txt")
+    assert str(ex.value) == "foo.txt"
+
+    with reader.open_resource("foo.txt") as f:
+        assert f.read() == b"bar"
+    with pytest.raises(FileNotFoundError) as ex:
+        reader.open_resource("notfound.txt")
+    assert str(ex.value) == "resfs/file/library/python/runtime_py3/test/resources/notfound.txt"
+
+    assert tuple(reader.contents()) == ("foo.txt", "data", "submodule")
diff --git a/library/python/runtime_py3/test/ya.make b/library/python/runtime_py3/test/ya.make
index fde64236dca..7e466d0a708 100644
--- a/library/python/runtime_py3/test/ya.make
+++ b/library/python/runtime_py3/test/ya.make
@@ -2,7 +2,9 @@ PY3TEST()
 
 STYLE_PYTHON()
 
-DEPENDS(library/python/runtime_py3/test/traceback)
+DEPENDS(
+    library/python/runtime_py3/test/traceback
+)
 
 PEERDIR(
     contrib/python/parameterized
@@ -29,6 +31,7 @@ RESOURCE_FILES(
     .dist-info/entry_points.txt
     .dist-info/top_level.txt
     resources/foo.txt
+    resources/data/my_data
     resources/submodule/bar.txt
 )