Bug 1803984 - Add tests for the interaction between speculative preloading and module...
[gecko.git] / build / mach_initialize.py
blobe28068800966d37d771bb366df6f29ae97a8bf5b
1 # This Source Code Form is subject to the terms of the Mozilla Public
2 # License, v. 2.0. If a copy of the MPL was not distributed with this
3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
5 import math
6 import os
7 import shutil
8 import sys
9 from pathlib import Path
11 if sys.version_info[0] < 3:
12 import __builtin__ as builtins
14 class MetaPathFinder(object):
15 pass
18 else:
19 from importlib.abc import MetaPathFinder
21 from types import ModuleType
23 STATE_DIR_FIRST_RUN = """
24 Mach and the build system store shared state in a common directory
25 on the filesystem. The following directory will be created:
29 If you would like to use a different directory, rename or move it to your
30 desired location, and set the MOZBUILD_STATE_PATH environment variable
31 accordingly.
32 """.strip()
35 CATEGORIES = {
36 "build": {
37 "short": "Build Commands",
38 "long": "Interact with the build system",
39 "priority": 80,
41 "post-build": {
42 "short": "Post-build Commands",
43 "long": "Common actions performed after completing a build.",
44 "priority": 70,
46 "testing": {
47 "short": "Testing",
48 "long": "Run tests.",
49 "priority": 60,
51 "ci": {
52 "short": "CI",
53 "long": "Taskcluster commands",
54 "priority": 59,
56 "devenv": {
57 "short": "Development Environment",
58 "long": "Set up and configure your development environment.",
59 "priority": 50,
61 "build-dev": {
62 "short": "Low-level Build System Interaction",
63 "long": "Interact with specific parts of the build system.",
64 "priority": 20,
66 "misc": {
67 "short": "Potpourri",
68 "long": "Potent potables and assorted snacks.",
69 "priority": 10,
71 "release": {
72 "short": "Release automation",
73 "long": "Commands for used in release automation.",
74 "priority": 5,
76 "disabled": {
77 "short": "Disabled",
78 "long": "The disabled commands are hidden by default. Use -v to display them. "
79 "These commands are unavailable for your current context, "
80 'run "mach <command>" to see why.',
81 "priority": 0,
86 def _activate_python_environment(topsrcdir, get_state_dir):
87 from mach.site import MachSiteManager
89 mach_environment = MachSiteManager.from_environment(
90 topsrcdir,
91 get_state_dir,
93 mach_environment.activate()
96 def _maybe_activate_mozillabuild_environment():
97 if sys.platform != "win32":
98 return
100 mozillabuild = Path(os.environ.get("MOZILLABUILD", r"C:\mozilla-build"))
101 os.environ.setdefault("MOZILLABUILD", str(mozillabuild))
102 assert mozillabuild.exists(), (
103 f'MozillaBuild was not found at "{mozillabuild}".\n'
104 "If it's installed in a different location, please "
105 'set the "MOZILLABUILD" environment variable '
106 "accordingly."
109 use_msys2 = (mozillabuild / "msys2").exists()
110 if use_msys2:
111 mozillabuild_msys_tools_path = mozillabuild / "msys2" / "usr" / "bin"
112 else:
113 mozillabuild_msys_tools_path = mozillabuild / "msys" / "bin"
115 paths_to_add = [mozillabuild_msys_tools_path, mozillabuild / "bin"]
116 existing_paths = [Path(p) for p in os.environ.get("PATH", "").split(os.pathsep)]
117 for new_path in paths_to_add:
118 if new_path not in existing_paths:
119 os.environ["PATH"] += f"{os.pathsep}{new_path}"
122 def initialize(topsrcdir):
123 # This directory was deleted in bug 1666345, but there may be some ignored
124 # files here. We can safely just delete it for the user so they don't have
125 # to clean the repo themselves.
126 deleted_dir = os.path.join(topsrcdir, "third_party", "python", "psutil")
127 if os.path.exists(deleted_dir):
128 shutil.rmtree(deleted_dir, ignore_errors=True)
130 # We need the "mach" module to access the logic to parse virtualenv
131 # requirements. Since that depends on "packaging" (and, transitively,
132 # "pyparsing"), we add those to the path too.
133 sys.path[0:0] = [
134 os.path.join(topsrcdir, module)
135 for module in (
136 os.path.join("python", "mach"),
137 os.path.join("third_party", "python", "packaging"),
138 os.path.join("third_party", "python", "pyparsing"),
142 from mach.util import get_state_dir, setenv
144 state_dir = _create_state_dir()
146 # normpath state_dir to normalize msys-style slashes.
147 _activate_python_environment(
148 topsrcdir, lambda: os.path.normpath(get_state_dir(True, topsrcdir=topsrcdir))
150 _maybe_activate_mozillabuild_environment()
152 import mach.base
153 import mach.main
154 from mach.main import MachCommandReference
156 # Centralized registry of available mach commands
157 MACH_COMMANDS = {
158 "valgrind-test": MachCommandReference("build/valgrind/mach_commands.py"),
159 "devtools-css-db": MachCommandReference(
160 "devtools/shared/css/generated/mach_commands.py"
162 "webidl-example": MachCommandReference("dom/bindings/mach_commands.py"),
163 "webidl-parser-test": MachCommandReference("dom/bindings/mach_commands.py"),
164 "hazards": MachCommandReference(
165 "js/src/devtools/rootAnalysis/mach_commands.py"
167 "reftest": MachCommandReference("layout/tools/reftest/mach_commands.py"),
168 "jstestbrowser": MachCommandReference("layout/tools/reftest/mach_commands.py"),
169 "crashtest": MachCommandReference("layout/tools/reftest/mach_commands.py"),
170 "android": MachCommandReference("mobile/android/mach_commands.py"),
171 "gradle": MachCommandReference("mobile/android/mach_commands.py"),
172 "gradle-install": MachCommandReference("mobile/android/mach_commands.py"),
173 "mach-commands": MachCommandReference(
174 "python/mach/mach/commands/commandinfo.py"
176 "mach-debug-commands": MachCommandReference(
177 "python/mach/mach/commands/commandinfo.py"
179 "mach-completion": MachCommandReference(
180 "python/mach/mach/commands/commandinfo.py"
182 "settings": MachCommandReference("python/mach/mach/commands/settings.py"),
183 "python": MachCommandReference("python/mach_commands.py"),
184 "python-test": MachCommandReference("python/mach_commands.py"),
185 "bootstrap": MachCommandReference("python/mozboot/mozboot/mach_commands.py"),
186 "vcs-setup": MachCommandReference("python/mozboot/mozboot/mach_commands.py"),
187 "artifact": MachCommandReference(
188 "python/mozbuild/mozbuild/artifact_commands.py"
190 "ide": MachCommandReference(
191 "python/mozbuild/mozbuild/backend/mach_commands.py"
193 "build": MachCommandReference("python/mozbuild/mozbuild/build_commands.py"),
194 "configure": MachCommandReference("python/mozbuild/mozbuild/build_commands.py"),
195 "resource-usage": MachCommandReference(
196 "python/mozbuild/mozbuild/build_commands.py"
198 "build-backend": MachCommandReference(
199 "python/mozbuild/mozbuild/build_commands.py"
201 "clang-tidy": MachCommandReference(
202 "python/mozbuild/mozbuild/code_analysis/mach_commands.py"
204 "static-analysis": MachCommandReference(
205 "python/mozbuild/mozbuild/code_analysis/mach_commands.py"
207 "prettier-format": MachCommandReference(
208 "python/mozbuild/mozbuild/code_analysis/mach_commands.py"
210 "clang-format": MachCommandReference(
211 "python/mozbuild/mozbuild/code_analysis/mach_commands.py"
213 "compileflags": MachCommandReference(
214 "python/mozbuild/mozbuild/compilation/codecomplete.py"
216 "mozbuild-reference": MachCommandReference(
217 "python/mozbuild/mozbuild/frontend/mach_commands.py"
219 "file-info": MachCommandReference(
220 "python/mozbuild/mozbuild/frontend/mach_commands.py"
222 "vendor": MachCommandReference(
223 "python/mozbuild/mozbuild/vendor/mach_commands.py"
225 "watch": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
226 "cargo": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
227 "doctor": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
228 "clobber": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
229 "show-log": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
230 "warnings-summary": MachCommandReference(
231 "python/mozbuild/mozbuild/mach_commands.py"
233 "warnings-list": MachCommandReference(
234 "python/mozbuild/mozbuild/mach_commands.py"
236 "gtest": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
237 "package": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
238 "install": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
239 "run": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
240 "buildsymbols": MachCommandReference(
241 "python/mozbuild/mozbuild/mach_commands.py"
243 "environment": MachCommandReference(
244 "python/mozbuild/mozbuild/mach_commands.py"
246 "repackage": MachCommandReference("python/mozbuild/mozbuild/mach_commands.py"),
247 "package-multi-locale": MachCommandReference(
248 "python/mozbuild/mozbuild/mach_commands.py"
250 "perftest": MachCommandReference(
251 "python/mozperftest/mozperftest/mach_commands.py"
253 "perftest-test": MachCommandReference(
254 "python/mozperftest/mozperftest/mach_commands.py"
256 "release": MachCommandReference(
257 "python/mozrelease/mozrelease/mach_commands.py"
259 "remote": MachCommandReference("remote/mach_commands.py"),
260 "puppeteer-test": MachCommandReference("remote/mach_commands.py"),
261 "generate-test-certs": MachCommandReference(
262 "security/manager/tools/mach_commands.py"
264 "taskgraph": MachCommandReference("taskcluster/mach_commands.py"),
265 "taskcluster-load-image": MachCommandReference("taskcluster/mach_commands.py"),
266 "taskcluster-build-image": MachCommandReference("taskcluster/mach_commands.py"),
267 "taskcluster-image-digest": MachCommandReference(
268 "taskcluster/mach_commands.py"
270 "release-history": MachCommandReference("taskcluster/mach_commands.py"),
271 "awsy-test": MachCommandReference("testing/awsy/mach_commands.py"),
272 "fetch-condprofile": MachCommandReference(
273 "testing/condprofile/mach_commands.py"
275 "run-condprofile": MachCommandReference("testing/condprofile/mach_commands.py"),
276 "firefox-ui-functional": MachCommandReference(
277 "testing/firefox-ui/mach_commands.py"
279 "geckodriver": MachCommandReference("testing/geckodriver/mach_commands.py"),
280 "addtest": MachCommandReference("testing/mach_commands.py"),
281 "test": MachCommandReference("testing/mach_commands.py"),
282 "cppunittest": MachCommandReference("testing/mach_commands.py"),
283 "jstests": MachCommandReference("testing/mach_commands.py"),
284 "jit-test": MachCommandReference("testing/mach_commands.py"),
285 "jsapi-tests": MachCommandReference("testing/mach_commands.py"),
286 "jsshell-bench": MachCommandReference("testing/mach_commands.py"),
287 "cramtest": MachCommandReference("testing/mach_commands.py"),
288 "test-info": MachCommandReference("testing/mach_commands.py"),
289 "rusttests": MachCommandReference("testing/mach_commands.py"),
290 "fluent-migration-test": MachCommandReference("testing/mach_commands.py"),
291 "marionette-test": MachCommandReference("testing/marionette/mach_commands.py"),
292 "mochitest": MachCommandReference("testing/mochitest/mach_commands.py"),
293 "geckoview-junit": MachCommandReference("testing/mochitest/mach_commands.py"),
294 "mozharness": MachCommandReference("testing/mozharness/mach_commands.py"),
295 "raptor": MachCommandReference("testing/raptor/mach_commands.py"),
296 "raptor-test": MachCommandReference("testing/raptor/mach_commands.py"),
297 "talos-test": MachCommandReference("testing/talos/mach_commands.py"),
298 "tps-build": MachCommandReference("testing/tps/mach_commands.py"),
299 "web-platform-tests": MachCommandReference(
300 "testing/web-platform/mach_commands.py"
302 "wpt": MachCommandReference("testing/web-platform/mach_commands.py"),
303 "web-platform-tests-update": MachCommandReference(
304 "testing/web-platform/mach_commands.py"
306 "wpt-update": MachCommandReference("testing/web-platform/mach_commands.py"),
307 "wpt-manifest-update": MachCommandReference(
308 "testing/web-platform/mach_commands.py"
310 "wpt-serve": MachCommandReference("testing/web-platform/mach_commands.py"),
311 "wpt-metadata-summary": MachCommandReference(
312 "testing/web-platform/mach_commands.py"
314 "wpt-metadata-merge": MachCommandReference(
315 "testing/web-platform/mach_commands.py"
317 "wpt-unittest": MachCommandReference("testing/web-platform/mach_commands.py"),
318 "wpt-test-paths": MachCommandReference("testing/web-platform/mach_commands.py"),
319 "wpt-fission-regressions": MachCommandReference(
320 "testing/web-platform/mach_commands.py"
322 "xpcshell-test": MachCommandReference("testing/xpcshell/mach_commands.py"),
323 "telemetry-tests-client": MachCommandReference(
324 "toolkit/components/telemetry/tests/marionette/mach_commands.py"
326 "data-review": MachCommandReference(
327 "toolkit/components/glean/build_scripts/mach_commands.py"
329 "perf-data-review": MachCommandReference(
330 "toolkit/components/glean/build_scripts/mach_commands.py"
332 "update-glean-tags": MachCommandReference(
333 "toolkit/components/glean/build_scripts/mach_commands.py"
335 "update-glean": MachCommandReference(
336 "toolkit/components/glean/build_scripts/mach_commands.py"
338 "browsertime": MachCommandReference("tools/browsertime/mach_commands.py"),
339 "compare-locales": MachCommandReference(
340 "tools/compare-locales/mach_commands.py"
342 "l10n-cross-channel": MachCommandReference("tools/lint/mach_commands.py"),
343 "busted": MachCommandReference("tools/mach_commands.py"),
344 "pastebin": MachCommandReference("tools/mach_commands.py"),
345 "mozregression": MachCommandReference("tools/mach_commands.py"),
346 "node": MachCommandReference("tools/mach_commands.py"),
347 "npm": MachCommandReference("tools/mach_commands.py"),
348 "logspam": MachCommandReference("tools/mach_commands.py"),
349 "doc": MachCommandReference("tools/moztreedocs/mach_commands.py"),
350 "install-moz-phab": MachCommandReference("tools/phabricator/mach_commands.py"),
351 "power": MachCommandReference("tools/power/mach_commands.py"),
352 "try": MachCommandReference("tools/tryselect/mach_commands.py"),
353 "import-pr": MachCommandReference("tools/vcs/mach_commands.py"),
354 "test-interventions": MachCommandReference(
355 "testing/webcompat/mach_commands.py"
357 "esmify": MachCommandReference("tools/esmify/mach_commands.py"),
358 "xpcshell": MachCommandReference("js/xpconnect/mach_commands.py"),
359 "uniffi": MachCommandReference(
360 "toolkit/components/uniffi-bindgen-gecko-js/mach_commands.py"
362 "storybook": MachCommandReference(
363 "browser/components/storybook/mach_commands.py"
365 "widgets": MachCommandReference("toolkit/content/widgets/mach_commands.py"),
368 # Set a reasonable limit to the number of open files.
370 # Some linux systems set `ulimit -n` to a very high number, which works
371 # well for systems that run servers, but this setting causes performance
372 # problems when programs close file descriptors before forking, like
373 # Python's `subprocess.Popen(..., close_fds=True)` (close_fds=True is the
374 # default in Python 3), or Rust's stdlib. In some cases, Firefox does the
375 # same thing when spawning processes. We would prefer to lower this limit
376 # to avoid such performance problems; processes spawned by `mach` will
377 # inherit the limit set here.
379 # The Firefox build defaults the soft limit to 1024, except for builds that
380 # do LTO, where the soft limit is 8192. We're going to default to the
381 # latter, since people do occasionally do LTO builds on their local
382 # machines, and requiring them to discover another magical setting after
383 # setting up an LTO build in the first place doesn't seem good.
385 # This code mimics the code in taskcluster/scripts/run-task.
386 try:
387 import resource
389 # Keep the hard limit the same, though, allowing processes to change
390 # their soft limit if they need to (Firefox does, for instance).
391 (soft, hard) = resource.getrlimit(resource.RLIMIT_NOFILE)
392 # Permit people to override our default limit if necessary via
393 # MOZ_LIMIT_NOFILE, which is the same variable `run-task` uses.
394 limit = os.environ.get("MOZ_LIMIT_NOFILE")
395 if limit:
396 limit = int(limit)
397 else:
398 # If no explicit limit is given, use our default if it's less than
399 # the current soft limit. For instance, the default on macOS is
400 # 256, so we'd pick that rather than our default.
401 limit = min(soft, 8192)
402 # Now apply the limit, if it's different from the original one.
403 if limit != soft:
404 resource.setrlimit(resource.RLIMIT_NOFILE, (limit, hard))
405 except ImportError:
406 # The resource module is UNIX only.
407 pass
409 def resolve_repository():
410 import mozversioncontrol
412 try:
413 # This API doesn't respect the vcs binary choices from configure.
414 # If we ever need to use the VCS binary here, consider something
415 # more robust.
416 return mozversioncontrol.get_repository_object(path=topsrcdir)
417 except (mozversioncontrol.InvalidRepoPath, mozversioncontrol.MissingVCSTool):
418 return None
420 def pre_dispatch_handler(context, handler, args):
421 # If --disable-tests flag was enabled in the mozconfig used to compile
422 # the build, tests will be disabled. Instead of trying to run
423 # nonexistent tests then reporting a failure, this will prevent mach
424 # from progressing beyond this point.
425 if handler.category == "testing" and not handler.ok_if_tests_disabled:
426 from mozbuild.base import BuildEnvironmentNotFoundException
428 try:
429 from mozbuild.base import MozbuildObject
431 # all environments should have an instance of build object.
432 build = MozbuildObject.from_environment()
433 if build is not None and not getattr(
434 build, "substs", {"ENABLE_TESTS": True}
435 ).get("ENABLE_TESTS"):
436 print(
437 "Tests have been disabled with --disable-tests.\n"
438 + "Remove the flag, and re-compile to enable tests."
440 sys.exit(1)
441 except BuildEnvironmentNotFoundException:
442 # likely automation environment, so do nothing.
443 pass
445 def post_dispatch_handler(
446 context, handler, instance, success, start_time, end_time, depth, args
448 """Perform global operations after command dispatch.
451 For now, we will use this to handle build system telemetry.
454 # Don't finalize telemetry data if this mach command was invoked as part of
455 # another mach command.
456 if depth != 1:
457 return
459 _finalize_telemetry_glean(
460 context.telemetry, handler.name == "bootstrap", success
463 def populate_context(key=None):
464 if key is None:
465 return
466 if key == "state_dir":
467 return state_dir
469 if key == "local_state_dir":
470 return get_state_dir(specific_to_topsrcdir=True)
472 if key == "topdir":
473 return topsrcdir
475 if key == "pre_dispatch_handler":
476 return pre_dispatch_handler
478 if key == "post_dispatch_handler":
479 return post_dispatch_handler
481 if key == "repository":
482 return resolve_repository()
484 raise AttributeError(key)
486 # Note which process is top-level so that recursive mach invocations can avoid writing
487 # telemetry data.
488 if "MACH_MAIN_PID" not in os.environ:
489 setenv("MACH_MAIN_PID", str(os.getpid()))
491 driver = mach.main.Mach(os.getcwd())
492 driver.populate_context_handler = populate_context
494 if not driver.settings_paths:
495 # default global machrc location
496 driver.settings_paths.append(state_dir)
497 # always load local repository configuration
498 driver.settings_paths.append(topsrcdir)
500 for category, meta in CATEGORIES.items():
501 driver.define_category(category, meta["short"], meta["long"], meta["priority"])
503 # Sparse checkouts may not have all mach_commands.py files. Ignore
504 # errors from missing files. Same for spidermonkey tarballs.
505 repo = resolve_repository()
506 if repo != "SOURCE":
507 missing_ok = (
508 repo is not None and repo.sparse_checkout_present()
509 ) or os.path.exists(os.path.join(topsrcdir, "INSTALL"))
510 else:
511 missing_ok = ()
512 driver.load_commands_from_spec(MACH_COMMANDS, topsrcdir, missing_ok=missing_ok)
513 return driver
516 def _finalize_telemetry_glean(telemetry, is_bootstrap, success):
517 """Submit telemetry collected by Glean.
519 Finalizes some metrics (command success state and duration, system information) and
520 requests Glean to send the collected data.
523 from mach.telemetry import MACH_METRICS_PATH
524 from mozbuild.telemetry import (
525 get_cpu_brand,
526 get_distro_and_version,
527 get_psutil_stats,
528 get_shell_info,
529 get_vscode_running,
532 mach_metrics = telemetry.metrics(MACH_METRICS_PATH)
533 mach_metrics.mach.duration.stop()
534 mach_metrics.mach.success.set(success)
535 system_metrics = mach_metrics.mach.system
536 cpu_brand = get_cpu_brand()
537 if cpu_brand:
538 system_metrics.cpu_brand.set(cpu_brand)
539 distro, version = get_distro_and_version()
540 system_metrics.distro.set(distro)
541 system_metrics.distro_version.set(version)
543 vscode_terminal, ssh_connection = get_shell_info()
544 system_metrics.vscode_terminal.set(vscode_terminal)
545 system_metrics.ssh_connection.set(ssh_connection)
546 system_metrics.vscode_running.set(get_vscode_running())
548 has_psutil, logical_cores, physical_cores, memory_total = get_psutil_stats()
549 if has_psutil:
550 # psutil may not be available (we may not have been able to download
551 # a wheel or build it from source).
552 system_metrics.logical_cores.add(logical_cores)
553 if physical_cores is not None:
554 system_metrics.physical_cores.add(physical_cores)
555 if memory_total is not None:
556 system_metrics.memory.accumulate(
557 int(math.ceil(float(memory_total) / (1024 * 1024 * 1024)))
559 telemetry.submit(is_bootstrap)
562 def _create_state_dir():
563 # Global build system and mach state is stored in a central directory. By
564 # default, this is ~/.mozbuild. However, it can be defined via an
565 # environment variable. We detect first run (by lack of this directory
566 # existing) and notify the user that it will be created. The logic for
567 # creation is much simpler for the "advanced" environment variable use
568 # case. For default behavior, we educate users and give them an opportunity
569 # to react.
570 state_dir = os.environ.get("MOZBUILD_STATE_PATH")
571 if state_dir:
572 if not os.path.exists(state_dir):
573 print(
574 "Creating global state directory from environment variable: {}".format(
575 state_dir
578 else:
579 state_dir = os.path.expanduser("~/.mozbuild")
580 if not os.path.exists(state_dir):
581 if not os.environ.get("MOZ_AUTOMATION"):
582 print(STATE_DIR_FIRST_RUN.format(state_dir))
584 print("Creating default state directory: {}".format(state_dir))
586 os.makedirs(state_dir, mode=0o770, exist_ok=True)
587 return state_dir
590 # Hook import such that .pyc/.pyo files without a corresponding .py file in
591 # the source directory are essentially ignored. See further below for details
592 # and caveats.
593 # Objdirs outside the source directory are ignored because in most cases, if
594 # a .pyc/.pyo file exists there, a .py file will be next to it anyways.
595 class ImportHook(object):
596 def __init__(self, original_import):
597 self._original_import = original_import
598 # Assume the source directory is the parent directory of the one
599 # containing this file.
600 self._source_dir = (
601 os.path.normcase(
602 os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
604 + os.sep
606 self._modules = set()
608 def __call__(self, name, globals=None, locals=None, fromlist=None, level=-1):
609 if sys.version_info[0] >= 3 and level < 0:
610 level = 0
612 # name might be a relative import. Instead of figuring out what that
613 # resolves to, which is complex, just rely on the real import.
614 # Since we don't know the full module name, we can't check sys.modules,
615 # so we need to keep track of which modules we've already seen to avoid
616 # to stat() them again when they are imported multiple times.
617 module = self._original_import(name, globals, locals, fromlist, level)
619 # Some tests replace modules in sys.modules with non-module instances.
620 if not isinstance(module, ModuleType):
621 return module
623 resolved_name = module.__name__
624 if resolved_name in self._modules:
625 return module
626 self._modules.add(resolved_name)
628 # Builtin modules don't have a __file__ attribute.
629 if not getattr(module, "__file__", None):
630 return module
632 # Note: module.__file__ is not always absolute.
633 path = os.path.normcase(os.path.abspath(module.__file__))
634 # Note: we could avoid normcase and abspath above for non pyc/pyo
635 # files, but those are actually rare, so it doesn't really matter.
636 if not path.endswith((".pyc", ".pyo")):
637 return module
639 # Ignore modules outside our source directory
640 if not path.startswith(self._source_dir):
641 return module
643 # If there is no .py corresponding to the .pyc/.pyo module we're
644 # loading, remove the .pyc/.pyo file, and reload the module.
645 # Since we already loaded the .pyc/.pyo module, if it had side
646 # effects, they will have happened already, and loading the module
647 # with the same name, from another directory may have the same side
648 # effects (or different ones). We assume it's not a problem for the
649 # python modules under our source directory (either because it
650 # doesn't happen or because it doesn't matter).
651 if not os.path.exists(module.__file__[:-1]):
652 if os.path.exists(module.__file__):
653 os.remove(module.__file__)
654 del sys.modules[module.__name__]
655 module = self(name, globals, locals, fromlist, level)
657 return module
660 # Hook import such that .pyc/.pyo files without a corresponding .py file in
661 # the source directory are essentially ignored. See further below for details
662 # and caveats.
663 # Objdirs outside the source directory are ignored because in most cases, if
664 # a .pyc/.pyo file exists there, a .py file will be next to it anyways.
665 class FinderHook(MetaPathFinder):
666 def __init__(self, klass):
667 # Assume the source directory is the parent directory of the one
668 # containing this file.
669 self._source_dir = (
670 os.path.normcase(
671 os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
673 + os.sep
675 self.finder_class = klass
677 def find_spec(self, full_name, paths=None, target=None):
678 spec = self.finder_class.find_spec(full_name, paths, target)
680 # Some modules don't have an origin.
681 if spec is None or spec.origin is None:
682 return spec
684 # Normalize the origin path.
685 path = os.path.normcase(os.path.abspath(spec.origin))
686 # Note: we could avoid normcase and abspath above for non pyc/pyo
687 # files, but those are actually rare, so it doesn't really matter.
688 if not path.endswith((".pyc", ".pyo")):
689 return spec
691 # Ignore modules outside our source directory
692 if not path.startswith(self._source_dir):
693 return spec
695 # If there is no .py corresponding to the .pyc/.pyo module we're
696 # resolving, remove the .pyc/.pyo file, and try again.
697 if not os.path.exists(spec.origin[:-1]):
698 if os.path.exists(spec.origin):
699 os.remove(spec.origin)
700 spec = self.finder_class.find_spec(full_name, paths, target)
702 return spec
705 # Additional hook for python >= 3.8's importlib.metadata.
706 class MetadataHook(FinderHook):
707 def find_distributions(self, *args, **kwargs):
708 return self.finder_class.find_distributions(*args, **kwargs)
711 def hook(finder):
712 has_find_spec = hasattr(finder, "find_spec")
713 has_find_distributions = hasattr(finder, "find_distributions")
714 if has_find_spec and has_find_distributions:
715 return MetadataHook(finder)
716 elif has_find_spec:
717 return FinderHook(finder)
718 return finder
721 # Install our hook. This can be deleted when the Python 3 migration is complete.
722 if sys.version_info[0] < 3:
723 builtins.__import__ = ImportHook(builtins.__import__)
724 else:
725 sys.meta_path = [hook(c) for c in sys.meta_path]