First implementation is `importlib.resources.files`

2 files changed, 82 insertions, 15 deletions

diff --git a/library/python/runtime_py3/importer.pxi b/library/python/runtime_py3/importer.pxi
index 56c5e439ec..3853678939 100644
--- a/library/python/runtime_py3/importer.pxi
+++ b/library/python/runtime_py3/importer.pxi
@@ -362,25 +362,19 @@ class ResourceImporter(object):
                 yield m
 
     def get_resource_reader(self, fullname):
-        try:
-            if not self.is_package(fullname):
-                return None
-        except ImportError:
-            return None
-        return _ResfsResourceReader(self, fullname)
+        import os
+        path = os.path.dirname(self.get_filename(fullname))
+        return _ResfsResourceReader(self, path)
 
 
 class _ResfsResourceReader:
 
-    def __init__(self, importer, fullname):
+    def __init__(self, importer, path):
         self.importer = importer
-        self.fullname = fullname
-
-        import os
-        self.prefix = "{}/".format(os.path.dirname(self.importer.get_filename(self.fullname)))
+        self.path = path
 
     def open_resource(self, resource):
-        path = f'{self.prefix}{resource}'
+        path = f'{self.path}/{resource}'
         from io import BytesIO
         try:
             return BytesIO(self.importer.get_data(path))
@@ -394,7 +388,7 @@ class _ResfsResourceReader:
         raise FileNotFoundError
 
     def is_resource(self, name):
-        path = f'{self.prefix}{name}'
+        path = f'{self.path}/{name}'
         try:
             self.importer.get_data(path)
         except OSError:
@@ -403,8 +397,9 @@ class _ResfsResourceReader:
 
     def contents(self):
         subdirs_seen = set()
-        for key in resfs_files(self.prefix):
-            relative = key[len(self.prefix):]
+        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)
@@ -412,6 +407,10 @@ class _ResfsResourceReader:
                 subdirs_seen.add(res_or_subdir)
                 yield _s(res_or_subdir)
 
+    def files(self):
+        import sitecustomize
+        return sitecustomize.ArcadiaResourceContainer(f"resfs/file/{self.path}/")
+
 
 class BuiltinSubmoduleImporter(BuiltinImporter):
     @classmethod
diff --git a/library/python/runtime_py3/sitecustomize.pyx b/library/python/runtime_py3/sitecustomize.pyx
index f1076b87d4..f566104c3c 100644
--- a/library/python/runtime_py3/sitecustomize.pyx
+++ b/library/python/runtime_py3/sitecustomize.pyx
@@ -1,4 +1,6 @@
 import pathlib
+import io
+import os
 import re
 import sys
 
@@ -6,12 +8,78 @@ import __res
 
 from importlib.abc import ResourceReader
 from importlib.metadata import Distribution, DistributionFinder, PackageNotFoundError, Prepared
+from importlib.resources.abc import Traversable
 
 ResourceReader.register(__res._ResfsResourceReader)
 
 METADATA_NAME = re.compile('^Name: (.*)$', re.MULTILINE)
 
 
+class ArcadiaResourceHandle(Traversable):
+    def __init__(self, key):
+        self.resfs_key = key
+
+    def is_file(self):
+        return True
+
+    def is_dir(self):
+        return False
+
+    def open(self, mode='r', *args, **kwargs):
+        data = __res.find(self.resfs_key.encode("utf-8"))
+        if data is None:
+            raise FileNotFoundError(self.resfs_key)
+
+        stream = io.BytesIO(data)
+
+        if 'b' not in mode:
+            stream = io.TextIOWrapper(stream, *args, **kwargs)
+
+        return stream
+
+    def joinpath(self, *name):
+        raise RuntimeError("Cannot traverse into a resource")
+
+    def iterdir(self):
+        return iter(())
+
+    @property
+    def name(self):
+        return os.path.basename(self.resfs_key)
+
+
+class ArcadiaResourceContainer(Traversable):
+    def __init__(self, prefix):
+        self.resfs_prefix = prefix
+
+    def is_dir(self):
+        return True
+
+    def is_file(self):
+        return False
+
+    def iterdir(self):
+        for key, path_without_prefix in __res.iter_keys(self.resfs_prefix.encode("utf-8")):
+            if b"/" in path_without_prefix:
+                name = path_without_prefix.decode("utf-8").split("/", maxsplit=1)[0]
+                yield ArcadiaResourceContainer(f"{self.resfs_prefix}{name}/")
+            else:
+                yield ArcadiaResourceHandle(key.decode("utf-8"))
+
+    def open(self, *args, **kwargs):
+        raise IsADirectoryError(self.resfs_prefix)
+
+    def joinpath(self, *descendants):
+        if not descendants:
+            return self
+
+        return ArcadiaResourceHandle(os.path.join(self.resfs_prefix, *descendants))
+
+    @property
+    def name(self):
+        return os.path.basename(self.resfs_prefix[:-1])
+
+
 class ArcadiaDistribution(Distribution):
 
     def __init__(self, prefix):