Coverage for src / vivarium / _compat / _compat.py: 62%
40 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-08 15:08 +0000
1"""Backward-compatible import redirects for the vivarium monorepo migration.
3Intercepts old-style imports (e.g. ``import layered_config_tree``,
4``from vivarium_public_health.disease import DiseaseModel``) and transparently
5redirects them to the new ``vivarium.*`` namespace, emitting a DeprecationWarning.
7Activated at interpreter startup via ``vivarium_compat.pth`` so the hook is in
8place before any user code runs.
10To add a redirect when a package migrates, add an entry to ``_REDIRECTS`` and
11bump the ``vivarium-compat`` version in ``pyproject.toml``.
13Remove this module once all downstream packages have released versions that use
14the new import paths and the deprecation period has ended.
15"""
17import importlib
18import importlib.abc
19import importlib.machinery
20import sys
21import warnings
22from types import ModuleType
24# Old import root -> new import root.
25# Activate an entry when its package has been migrated into the monorepo.
26# The hook will fail loudly if the new location does not yet exist, so do not
27# enable an entry before its target package is released.
28_REDIRECTS: dict[str, str] = {
29 # Renamed top-level packages (enable when each package migrates)
30 "vivarium_profiling": "vivarium.profiling",
31 # "vivarium_public_health": "vivarium.public_health",
32 # "vivarium_cluster_tools": "vivarium.cluster_tools",
33 # "vivarium_testing_utils": "vivarium.testing_utils",
34 # "vivarium.examples": "vivarium.core.examples",
35 # "vivarium.framework": "vivarium.core.framework",
36 # "vivarium.interface": "vivarium.core.interface",
37 # "vivarium.component": "vivarium.core.component",
38 # "vivarium.exceptions": "vivarium.core.exceptions",
39 # "vivarium.manager": "vivarium.core.manager",
40 # "vivarium.testing_utilities": "vivarium.core.testing_utilities",
41 # "vivarium.types": "vivarium.core.types",
42 # "vivarium_helpers": "vivarium.helpers",
43 # "risk_distributions": "vivarium.risk_distributions",
44 # "gbd_mapping": "vivarium.gbd_mapping",
45 # "layered_config_tree": "vivarium.config_tree",
46}
48# Tracks which old names are currently being resolved to prevent infinite
49# recursion if a redirect target somehow re-triggers the same old-name import.
50_resolving: set[str] = set()
53def _match(fullname: str) -> tuple[str, str] | None:
54 """Return (old_prefix, new_prefix) if fullname matches a redirect, else None.
56 Longer prefixes take precedence so more-specific entries win over broader ones.
57 """
58 for old, new in sorted(_REDIRECTS.items(), key=lambda x: -len(x[0])):
59 if fullname == old or fullname.startswith(old + "."):
60 return old, new
61 return None
64class _CompatFinder(importlib.abc.MetaPathFinder):
65 """Meta-path finder that redirects deprecated import paths to new locations."""
67 def find_spec(
68 self,
69 fullname: str,
70 path: object,
71 target: ModuleType | None = None,
72 ) -> importlib.machinery.ModuleSpec | None:
73 match = _match(fullname)
74 if match is None:
75 return None
77 old_prefix, new_prefix = match
78 new_name = new_prefix + fullname[len(old_prefix) :]
80 # Only warn on first import. After exec_module runs, sys.modules[fullname]
81 # is set to the real module, so subsequent imports return it directly without
82 # ever reaching find_spec again. This branch fires only in unusual re-entry
83 # scenarios (e.g. reload()).
84 if fullname not in sys.modules:
85 # stacklevel=2 points into importlib internals rather than the caller's
86 # import statement - the exact depth is CPython-version-dependent.
87 # The warning message itself is the actionable part.
88 warnings.warn(
89 f"'{fullname}' has moved to '{new_name}'. "
90 "Update your imports. This redirect will be removed in a future release.",
91 DeprecationWarning,
92 stacklevel=2,
93 )
94 # submodule_search_locations is intentionally not set on the spec. CPython's
95 # import machinery calls exec_module before checking __path__ on any child import,
96 # and exec_module replaces sys.modules[fullname] with the real module (which already
97 # has the correct __path__). Setting it to [] would be a no-op at best and misleading
98 # at worst since we don't know at spec-creation time whether the target is a package.
99 return importlib.machinery.ModuleSpec(
100 fullname, _CompatLoader(fullname, new_name)
101 )
104class _CompatLoader(importlib.abc.Loader):
105 """Loads the real module at the new location and aliases it under the old name."""
107 def __init__(self, old_name: str, new_name: str) -> None:
108 self._old_name = old_name
109 self._new_name = new_name
111 def exec_module(self, module: ModuleType) -> None:
112 if self._old_name in _resolving:
113 raise ImportError(
114 f"Circular redirect detected: '{self._old_name}' -> '{self._new_name}'"
115 )
116 _resolving.add(self._old_name)
117 try:
118 real = importlib.import_module(self._new_name)
119 # Register under the old name so subsequent imports hit sys.modules directly.
120 sys.modules[self._old_name] = real
121 # Defensive: covers the case where another import hook holds a direct
122 # reference to the placeholder module object rather than re-reading sys.modules.
123 module.__dict__.update(real.__dict__)
124 # __spec__.name will show the new name (e.g. "vivarium.config_tree"), not the
125 # old one. This is intentional: sys.modules[old_name] IS real, so its metadata
126 # correctly describes its actual location. Users inspecting __spec__ after
127 # migrating their imports will see the right thing.
128 module.__spec__ = real.__spec__
129 finally:
130 _resolving.discard(self._old_name)
133def install_compat_finder() -> None:
134 """Install the compat finder into sys.meta_path (idempotent)."""
135 if not any(isinstance(f, _CompatFinder) for f in sys.meta_path):
136 sys.meta_path.insert(0, _CompatFinder())