Coverage for src / vivarium_compat / _compat.py: 69%
52 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-21 21:38 +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``. Entries are safe
12to add *before* the target package is released: if the new module isn't
13importable, the hook falls back to letting the old on-disk package resolve
14normally, so installations during the transition window aren't broken.
16Remove this module once all downstream packages have released versions that use
17the new import paths and the deprecation period has ended.
18"""
20import importlib
21import importlib.abc
22import importlib.machinery
23import sys
24import warnings
25from types import ModuleType
27# Old import root -> new import root.
28# Entries here are safe to populate ahead of the target package's release:
29# if the new target isn't installed yet, the hook falls back to the old
30# package's normal on-disk location (so existing installations during the
31# transition window keep working).
32_REDIRECTS: dict[str, str] = {
33 # Renamed top-level packages
34 "vivarium_profiling": "vivarium.profiling",
35 "risk_distributions": "vivarium.risk_distributions",
36 "gbd_mapping": "vivarium.gbd_mapping",
37 "gbd_mapping_generator": "vivarium.gbd_mapping_generator",
38 "layered_config_tree": "vivarium.config_tree",
39 # vivarium -> vivarium-engine: top-level submodules moved under
40 # vivarium.engine.*; the entries below redirect submodule imports
41 # (e.g. ``from vivarium.framework.engine import Builder``).
42 # NOTE: Attribute imports off the bare namespace (e.g. ``from vivarium
43 # import Component``) are handled separately by the ``__getattr__``
44 # in vivarium-engine's ``vivarium/__init__.py``; the compat hook
45 # cannot intercept those without breaking sibling-namespace lookups.
46 "vivarium.examples": "vivarium.engine.examples",
47 "vivarium.framework": "vivarium.engine.framework",
48 "vivarium.interface": "vivarium.engine.interface",
49 "vivarium.component": "vivarium.engine.component",
50 "vivarium.exceptions": "vivarium.engine.exceptions",
51 "vivarium.manager": "vivarium.engine.manager",
52 "vivarium.testing_utilities": "vivarium.engine.testing_utilities",
53 "vivarium.types": "vivarium.engine.types",
54 # Not-yet-migrated libs; uncomment as each lands in the monorepo:
55 # "vivarium_public_health": "vivarium.public_health",
56 # "vivarium_cluster_tools": "vivarium.cluster_tools",
57 # "vivarium_testing_utils": "vivarium.testing_utils",
58 # "vivarium_helpers": "vivarium.helpers",
59}
61# Tracks which old names are currently being resolved to prevent infinite
62# recursion if a redirect target somehow re-triggers the same old-name import.
63_resolving: set[str] = set()
66def _match(fullname: str) -> tuple[str, str] | None:
67 """Return (old_prefix, new_prefix) if fullname matches a redirect, else None.
69 Longer prefixes take precedence so more-specific entries win over broader ones.
70 """
71 for old, new in sorted(_REDIRECTS.items(), key=lambda x: -len(x[0])):
72 if fullname == old or fullname.startswith(old + "."):
73 return old, new
74 return None
77class _CompatFinder(importlib.abc.MetaPathFinder):
78 """Meta-path finder that redirects deprecated import paths to new locations."""
80 def find_spec(
81 self,
82 fullname: str,
83 path: object,
84 target: ModuleType | None = None,
85 ) -> importlib.machinery.ModuleSpec | None:
86 match = _match(fullname)
87 if match is None:
88 return None
90 old_prefix, new_prefix = match
91 new_name = new_prefix + fullname[len(old_prefix) :]
93 # Only warn on first import. After exec_module runs, sys.modules[fullname]
94 # is set to the real module, so subsequent imports return it directly without
95 # ever reaching find_spec again. This branch fires only in unusual re-entry
96 # scenarios (e.g. reload()).
97 if fullname not in sys.modules:
98 # stacklevel=2 points into importlib internals rather than the caller's
99 # import statement - the exact depth is CPython-version-dependent.
100 # The warning message itself is the actionable part.
101 warnings.warn(
102 f"'{fullname}' has moved to '{new_name}'. "
103 "Update your imports. This redirect will be removed in a future release.",
104 DeprecationWarning,
105 stacklevel=2,
106 )
107 # submodule_search_locations is intentionally not set on the spec. CPython's
108 # import machinery calls exec_module before checking __path__ on any child import,
109 # and exec_module replaces sys.modules[fullname] with the real module (which already
110 # has the correct __path__). Setting it to [] would be a no-op at best and misleading
111 # at worst since we don't know at spec-creation time whether the target is a package.
112 return importlib.machinery.ModuleSpec(
113 fullname, _CompatLoader(fullname, new_name)
114 )
117class _CompatLoader(importlib.abc.Loader):
118 """Loads the real module at the new location and aliases it under the old name."""
120 def __init__(self, old_name: str, new_name: str) -> None:
121 self._old_name = old_name
122 self._new_name = new_name
124 def exec_module(self, module: ModuleType) -> None:
125 if self._old_name in _resolving:
126 raise ImportError(
127 f"Circular redirect detected: '{self._old_name}' -> '{self._new_name}'"
128 )
129 _resolving.add(self._old_name)
130 try:
131 try:
132 real = importlib.import_module(self._new_name)
133 except ModuleNotFoundError:
134 # FIXME: MIC-7100 Revert when all packages are migrated
135 # New target isn't installed. Fall back to the old name's actual
136 # on-disk location so installations during the transition window
137 # — old package still installed, new package not yet released —
138 # keep working. The DeprecationWarning was already emitted in
139 # find_spec(); users will see it whether or not the fallback fires.
140 real = _import_bypassing_compat(self._old_name)
141 # Register under the old name so subsequent imports hit sys.modules directly.
142 sys.modules[self._old_name] = real
143 # Defensive: covers the case where another import hook holds a direct
144 # reference to the placeholder module object rather than re-reading sys.modules.
145 module.__dict__.update(real.__dict__)
146 # __spec__.name will show the new name (e.g. "vivarium.config_tree"), not the
147 # old one. This is intentional: sys.modules[old_name] IS real, so its metadata
148 # correctly describes its actual location. Users inspecting __spec__ after
149 # migrating their imports will see the right thing.
150 module.__spec__ = real.__spec__
151 finally:
152 _resolving.discard(self._old_name)
155def _import_bypassing_compat(name: str) -> ModuleType:
156 """Import `name` without going through our compat finder.
158 Used to fall back to a name's real on-disk location when the redirect target
159 isn't installed. Temporarily removes our finder from sys.meta_path so the
160 standard PathFinder resolves the name directly. Also clears any placeholder
161 sys.modules entry the import system set when our find_spec returned a spec.
162 """
163 sys.modules.pop(name, None)
164 saved: list[tuple[int, _CompatFinder]] = [
165 (i, f) for i, f in enumerate(sys.meta_path) if isinstance(f, _CompatFinder)
166 ]
167 for _, f in saved:
168 sys.meta_path.remove(f)
169 try:
170 return importlib.import_module(name)
171 finally:
172 # Reinsert in original order so any other meta_path entries stay in place.
173 for i, f in saved:
174 sys.meta_path.insert(i, f)
177def install_compat_finder() -> None:
178 """Install the compat finder into sys.meta_path (idempotent)."""
179 if not any(isinstance(f, _CompatFinder) for f in sys.meta_path):
180 sys.meta_path.insert(0, _CompatFinder())