From 5aab00ecccf3629ae777e28c7dbc4b04fa914401 Mon Sep 17 00:00:00 2001
From: Mike Hommey <mh+mozilla@glandium.org>
Date: Tue, 4 May 2021 20:39:58 +0000
Subject: [PATCH] Bug 1708547 - Port the mach bootstrap import hook to python
 3. r=mhentges,firefox-build-system-reviewers

Differential Revision: https://phabricator.services.mozilla.com/D113857
---
 build/mach_bootstrap.py | 71 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 70 insertions(+), 1 deletion(-)

diff --git a/build/mach_bootstrap.py b/build/mach_bootstrap.py
index 4e68b974371e..1f01126154fe 100644
--- a/build/mach_bootstrap.py
+++ b/build/mach_bootstrap.py
@@ -13,8 +13,14 @@ import sys
 
 if sys.version_info[0] < 3:
     import __builtin__ as builtins
+
+    class MetaPathFinder(object):
+        pass
+
+
 else:
-    import builtins
+    from importlib.abc import MetaPathFinder
+
 
 from types import ModuleType
 
@@ -508,6 +514,69 @@ class ImportHook(object):
         return module
 
 
+# Hook import such that .pyc/.pyo files without a corresponding .py file in
+# the source directory are essentially ignored. See further below for details
+# and caveats.
+# Objdirs outside the source directory are ignored because in most cases, if
+# a .pyc/.pyo file exists there, a .py file will be next to it anyways.
+class FinderHook(MetaPathFinder):
+    def __init__(self, klass):
+        # Assume the source directory is the parent directory of the one
+        # containing this file.
+        self._source_dir = (
+            os.path.normcase(
+                os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
+            )
+            + os.sep
+        )
+        self.finder_class = klass
+
+    def find_spec(self, full_name, paths=None, target=None):
+        spec = self.finder_class.find_spec(full_name, paths, target)
+
+        # Some modules don't have an origin.
+        if spec is None or spec.origin is None:
+            return spec
+
+        # Normalize the origin path.
+        path = os.path.normcase(os.path.abspath(spec.origin))
+        # Note: we could avoid normcase and abspath above for non pyc/pyo
+        # files, but those are actually rare, so it doesn't really matter.
+        if not path.endswith((".pyc", ".pyo")):
+            return spec
+
+        # Ignore modules outside our source directory
+        if not path.startswith(self._source_dir):
+            return spec
+
+        # If there is no .py corresponding to the .pyc/.pyo module we're
+        # resolving, remove the .pyc/.pyo file, and try again.
+        if not os.path.exists(spec.origin[:-1]):
+            if os.path.exists(spec.origin):
+                os.remove(spec.origin)
+            spec = self.finder_class.find_spec(full_name, paths, target)
+
+        return spec
+
+
+# Additional hook for python >= 3.8's importlib.metadata.
+class MetadataHook(FinderHook):
+    def find_distributions(self, *args, **kwargs):
+        return self.finder_class.find_distributions(*args, **kwargs)
+
+
+def hook(finder):
+    has_find_spec = hasattr(finder, "find_spec")
+    has_find_distributions = hasattr(finder, "find_distributions")
+    if has_find_spec and has_find_distributions:
+        return MetadataHook(finder)
+    elif has_find_spec:
+        return FinderHook(finder)
+    return finder
+
+
 # Install our hook. This can be deleted when the Python 3 migration is complete.
 if sys.version_info[0] < 3:
     builtins.__import__ = ImportHook(builtins.__import__)
+else:
+    sys.meta_path = [hook(c) for c in sys.meta_path]