From b51e05fc7aa97e1105d6fc07f474173c9708c7d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bern=C3=A1t=20G=C3=A1bor?= <bgabor8@bloomberg.net>
Date: Thu, 17 Oct 2024 14:31:00 -0700
Subject: [PATCH] Fix broken Windows zipapp
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Bernát Gábor <bgabor8@bloomberg.net>
---
 tasks/__main__zipapp.py | 75 ++++++++++++++++++-----------------------
 1 file changed, 33 insertions(+), 42 deletions(-)

diff --git a/tasks/__main__zipapp.py b/tasks/__main__zipapp.py
index 01e3efaf2..cd6b478b1 100644
--- a/tasks/__main__zipapp.py
+++ b/tasks/__main__zipapp.py
@@ -4,9 +4,10 @@
 import os
 import sys
 import zipfile
+from importlib.abc import SourceLoader
+from importlib.util import spec_from_file_location
 
 ABS_HERE = os.path.abspath(os.path.dirname(__file__))
-NEW_IMPORT_SYSTEM = sys.version_info[0] >= 3  # noqa: PLR2004
 
 
 class VersionPlatformSelect:
@@ -69,18 +70,34 @@ def _register_distutils_finder(self):
         if "distlib" not in self.modules:
             return
 
+        class Resource:
+            def __init__(self, content) -> None:
+                self.bytes = content
+                self.is_container = False
+
         class DistlibFinder:
             def __init__(self, path, loader) -> None:
                 self.path = path
                 self.loader = loader
 
             def find(self, name):
-                class Resource:
-                    def __init__(self, content) -> None:
-                        self.bytes = content
-
-                full_path = os.path.join(self.path, name)
-                return Resource(self.loader.get_data(full_path))
+                return Resource(self.loader.get_data(os.path.join(self.path, name)))
+
+            def iterator(self, resource_name):
+                resource = self.find(resource_name)
+                if resource is not None:
+                    todo = [resource]
+                    while todo:
+                        resource = todo.pop(0)
+                        yield resource
+                        if resource.is_container:
+                            rname = resource.name
+                            for name in resource.resources:
+                                child = self.find(f"{rname}/{name}" if rname else name)
+                                if child.is_container:
+                                    todo.append(child)
+                                else:
+                                    yield child
 
         from distlib.resources import register_finder  # noqa: PLC0415
 
@@ -113,41 +130,15 @@ def locate_file(self, path):
     return _VER_DISTRIBUTION_CLASS
 
 
-if NEW_IMPORT_SYSTEM:
-    from importlib.abc import SourceLoader
-    from importlib.util import spec_from_file_location
-
-    class VersionedFindLoad(VersionPlatformSelect, SourceLoader):
-        def find_spec(self, fullname, path, target=None):  # noqa: ARG002
-            zip_path = self.find_mod(fullname)
-            if zip_path is not None:
-                return spec_from_file_location(name=fullname, loader=self)
-            return None
-
-        def module_repr(self, module):
-            raise NotImplementedError
-
-else:
-    from imp import new_module
-
-    class VersionedFindLoad(VersionPlatformSelect):
-        def find_module(self, fullname, path=None):  # noqa: ARG002
-            return self if self.find_mod(fullname) else None
-
-        def load_module(self, fullname):
-            filename = self.get_filename(fullname)
-            code = self.get_data(filename)
-            mod = sys.modules.setdefault(fullname, new_module(fullname))
-            mod.__file__ = filename
-            mod.__loader__ = self
-            is_package = filename.endswith("__init__.py")
-            if is_package:
-                mod.__path__ = [os.path.dirname(filename)]
-                mod.__package__ = fullname
-            else:
-                mod.__package__ = fullname.rpartition(".")[0]
-            exec(code, mod.__dict__)  # noqa: S102
-            return mod
+class VersionedFindLoad(VersionPlatformSelect, SourceLoader):
+    def find_spec(self, fullname, path, target=None):  # noqa: ARG002
+        zip_path = self.find_mod(fullname)
+        if zip_path is not None:
+            return spec_from_file_location(name=fullname, loader=self)
+        return None
+
+    def module_repr(self, module):
+        raise NotImplementedError
 
 
 def run():