Coverage for src / vivarium_compat / _compat.py: 69%

1"""Backward-compatible import redirects for the vivarium monorepo migration. 

2 

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. 

6 

7Activated at interpreter startup via ``vivarium_compat.pth`` so the hook is in 

8place before any user code runs. 

9 

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. 

15 

16Remove this module once all downstream packages have released versions that use 

17the new import paths and the deprecation period has ended. 

18""" 

19 

20import importlib 

21import importlib.abc 

22import importlib.machinery 

23import sys 

24import warnings 

25from types import ModuleType 

26 

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 # FIXME [MIC-7113]: the layered_config_tree redirect is disabled until vivarium 

39 # is migrated and all `from layered_config_tree import ...` statements are updated. 

40 # "layered_config_tree": "vivarium.config_tree", 

41 # "vivarium_public_health": "vivarium.public_health", 

42 # "vivarium_cluster_tools": "vivarium.cluster_tools", 

43 # "vivarium_testing_utils": "vivarium.testing_utils", 

44 # "vivarium.examples": "vivarium.core.examples", 

45 # "vivarium.framework": "vivarium.core.framework", 

46 # "vivarium.interface": "vivarium.core.interface", 

47 # "vivarium.component": "vivarium.core.component", 

48 # "vivarium.exceptions": "vivarium.core.exceptions", 

49 # "vivarium.manager": "vivarium.core.manager", 

50 # "vivarium.testing_utilities": "vivarium.core.testing_utilities", 

51 # "vivarium.types": "vivarium.core.types", 

52 # "vivarium_helpers": "vivarium.helpers", 

53} 

54 

55# Tracks which old names are currently being resolved to prevent infinite 

56# recursion if a redirect target somehow re-triggers the same old-name import. 

57_resolving: set[str] = set() 

58 

59 

60def _match(fullname: str) -> tuple[str, str] | None: 

61 """Return (old_prefix, new_prefix) if fullname matches a redirect, else None. 

62 

63 Longer prefixes take precedence so more-specific entries win over broader ones. 

64 """ 

65 for old, new in sorted(_REDIRECTS.items(), key=lambda x: -len(x[0])): 

66 if fullname == old or fullname.startswith(old + "."): 

67 return old, new 

68 return None 

69 

70 

71class _CompatFinder(importlib.abc.MetaPathFinder): 

72 """Meta-path finder that redirects deprecated import paths to new locations.""" 

73 

74 def find_spec( 

75 self, 

76 fullname: str, 

77 path: object, 

78 target: ModuleType | None = None, 

79 ) -> importlib.machinery.ModuleSpec | None: 

80 match = _match(fullname) 

81 if match is None: 

82 return None 

83 

84 old_prefix, new_prefix = match 

85 new_name = new_prefix + fullname[len(old_prefix) :] 

86 

87 # Only warn on first import. After exec_module runs, sys.modules[fullname] 

88 # is set to the real module, so subsequent imports return it directly without 

89 # ever reaching find_spec again. This branch fires only in unusual re-entry 

90 # scenarios (e.g. reload()). 

91 if fullname not in sys.modules: 

92 # stacklevel=2 points into importlib internals rather than the caller's 

93 # import statement - the exact depth is CPython-version-dependent. 

94 # The warning message itself is the actionable part. 

95 warnings.warn( 

96 f"'{fullname}' has moved to '{new_name}'. " 

97 "Update your imports. This redirect will be removed in a future release.", 

98 DeprecationWarning, 

99 stacklevel=2, 

100 ) 

101 # submodule_search_locations is intentionally not set on the spec. CPython's 

102 # import machinery calls exec_module before checking __path__ on any child import, 

103 # and exec_module replaces sys.modules[fullname] with the real module (which already 

104 # has the correct __path__). Setting it to [] would be a no-op at best and misleading 

105 # at worst since we don't know at spec-creation time whether the target is a package. 

106 return importlib.machinery.ModuleSpec( 

107 fullname, _CompatLoader(fullname, new_name) 

108 ) 

109 

110 

111class _CompatLoader(importlib.abc.Loader): 

112 """Loads the real module at the new location and aliases it under the old name.""" 

113 

114 def __init__(self, old_name: str, new_name: str) -> None: 

115 self._old_name = old_name 

116 self._new_name = new_name 

117 

118 def exec_module(self, module: ModuleType) -> None: 

119 if self._old_name in _resolving: 

120 raise ImportError( 

121 f"Circular redirect detected: '{self._old_name}' -> '{self._new_name}'" 

122 ) 

123 _resolving.add(self._old_name) 

124 try: 

125 try: 

126 real = importlib.import_module(self._new_name) 

127 except ModuleNotFoundError: 

128 # FIXME: MIC-7100 Revert when all packages are migrated 

129 # New target isn't installed. Fall back to the old name's actual 

130 # on-disk location so installations during the transition window 

131 # — old package still installed, new package not yet released — 

132 # keep working. The DeprecationWarning was already emitted in 

133 # find_spec(); users will see it whether or not the fallback fires. 

134 real = _import_bypassing_compat(self._old_name) 

135 # Register under the old name so subsequent imports hit sys.modules directly. 

136 sys.modules[self._old_name] = real 

137 # Defensive: covers the case where another import hook holds a direct 

138 # reference to the placeholder module object rather than re-reading sys.modules. 

139 module.__dict__.update(real.__dict__) 

140 # __spec__.name will show the new name (e.g. "vivarium.config_tree"), not the 

141 # old one. This is intentional: sys.modules[old_name] IS real, so its metadata 

142 # correctly describes its actual location. Users inspecting __spec__ after 

143 # migrating their imports will see the right thing. 

144 module.__spec__ = real.__spec__ 

145 finally: 

146 _resolving.discard(self._old_name) 

147 

148 

149def _import_bypassing_compat(name: str) -> ModuleType: 

150 """Import `name` without going through our compat finder. 

151 

152 Used to fall back to a name's real on-disk location when the redirect target 

153 isn't installed. Temporarily removes our finder from sys.meta_path so the 

154 standard PathFinder resolves the name directly. Also clears any placeholder 

155 sys.modules entry the import system set when our find_spec returned a spec. 

156 """ 

157 sys.modules.pop(name, None) 

158 saved: list[tuple[int, _CompatFinder]] = [ 

159 (i, f) for i, f in enumerate(sys.meta_path) if isinstance(f, _CompatFinder) 

160 ] 

161 for _, f in saved: 

162 sys.meta_path.remove(f) 

163 try: 

164 return importlib.import_module(name) 

165 finally: 

166 # Reinsert in original order so any other meta_path entries stay in place. 

167 for i, f in saved: 

168 sys.meta_path.insert(i, f) 

169 

170 

171def install_compat_finder() -> None: 

172 """Install the compat finder into sys.meta_path (idempotent).""" 

173 if not any(isinstance(f, _CompatFinder) for f in sys.meta_path): 

174 sys.meta_path.insert(0, _CompatFinder())