-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Expand file tree
/
Copy pathimportutils.py
More file actions
309 lines (251 loc) · 10 KB
/
importutils.py
File metadata and controls
309 lines (251 loc) · 10 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
"""
Module importing related utilities.
Author
* Sylvain de Langen 2024
"""
import importlib
import inspect
import os
import sys
import warnings
from types import ModuleType
from typing import List, Optional
class LazyModule(ModuleType):
"""Defines a module type that lazily imports the target module, thus
exposing contents without importing the target module needlessly.
Arguments
---------
name : str
Name of the module.
target : str
Module to be loading lazily.
package : str, optional
If specified, the target module load will be relative to this package.
Depending on how you inject the lazy module into the environment, you
may choose to specify the package here, or you may choose to include it
into the `name` with the dot syntax.
e.g. see how :func:`~lazy_export` and :func:`~deprecated_redirect`
differ.
"""
def __init__(
self,
name: str,
target: str,
package: Optional[str],
):
super().__init__(name)
self.target = target
self.lazy_module = None
self.package = package
def ensure_module(self, stacklevel: int) -> ModuleType:
"""Ensures that the target module is imported and available as
`self.lazy_module`, also returning it.
Arguments
---------
stacklevel : int
The stack trace level of the function that caused the import to
occur, relative to the **caller** of this function (e.g. if in
function `f` you call `ensure_module(1)`, it will refer to the
function that called `f`).
Raises
------
AttributeError
When the function responsible for the import attempt is found to be
`inspect.py`, we raise an `AttributeError` here. This is because
some code will inadvertently cause our modules to be imported, such
as some of PyTorch's op registering machinery.
Returns
-------
The target module after ensuring it is imported.
"""
importer_frame = None
# NOTE: ironically, calling this causes getframeinfo to call into
# `findsource` -> `getmodule` -> ourselves here
# bear that in mind if you are debugging and checking out the trace.
# also note that `_getframe` is an implementation detail, but it is
# somewhat non-critical to us.
try:
importer_frame = inspect.getframeinfo(sys._getframe(stacklevel + 1))
except AttributeError:
warnings.warn(
"Failed to inspect frame to check if we should ignore "
"importing a module lazily. This relies on a CPython "
"implementation detail, report an issue if you see this with "
"standard Python and include your version number."
)
if importer_frame is not None and importer_frame.filename.endswith(
"/inspect.py"
):
raise AttributeError()
if self.lazy_module is None:
try:
if self.package is None:
self.lazy_module = importlib.import_module(self.target)
else:
self.lazy_module = importlib.import_module(
f".{self.target}", self.package
)
except Exception as e:
raise ImportError(f"Lazy import of {repr(self)} failed") from e
return self.lazy_module
def __repr__(self) -> str:
return f"LazyModule(package={self.package}, target={self.target}, loaded={self.lazy_module is not None})"
def __getattr__(self, attr):
# NOTE: exceptions here get eaten and not displayed
return getattr(self.ensure_module(1), attr)
class DeprecatedModuleRedirect(LazyModule):
"""Defines a module type that lazily imports the target module using
:class:`~LazyModule`, but logging a deprecation warning when the import
is actually being performed.
This is only the module type itself; if you want to define a redirection,
use :func:`~deprecated_redirect` instead.
Arguments
---------
old_import : str
Old module import path e.g. `mypackage.myoldmodule`
new_import : str
New module import path e.g. `mypackage.mynewcoolmodule.mycoolsubmodule`
extra_reason : str, optional
If specified, extra text to attach to the warning for clarification
(e.g. justifying why the move has occurred, or additional problems to
look out for).
"""
def __init__(
self,
old_import: str,
new_import: str,
extra_reason: Optional[str] = None,
):
super().__init__(name=old_import, target=new_import, package=None)
self.old_import = old_import
self.extra_reason = extra_reason
def _redirection_warn(self):
"""Emits the warning for the redirection (with the extra reason if
provided)."""
warning_text = (
f"Module '{self.old_import}' was deprecated, redirecting to "
f"'{self.target}'. Please update your script."
)
if self.extra_reason is not None:
warning_text += f" {self.extra_reason}"
# NOTE: we are not using DeprecationWarning because this gets ignored by
# default, even though we consider the warning to be rather important
# in the context of SB
warnings.warn(
warning_text,
# category=DeprecationWarning,
stacklevel=4, # ensure_module <- __getattr__ <- python <- user
)
def ensure_module(self, stacklevel: int) -> ModuleType:
should_warn = self.lazy_module is None
# can fail with exception if the module shouldn't be imported, so only
# actually emit the warning later
module = super().ensure_module(stacklevel + 1)
if should_warn:
self._redirection_warn()
return module
def find_imports(file_path: str, find_subpackages: bool = False) -> List[str]:
"""Returns a list of importable scripts in the same module as the specified
file. e.g. if you have `foo/__init__.py` and `foo/bar.py`, then
`files_in_module("foo/__init__.py")` then the result will be `["bar"]`.
Not recursive; this is only applies to the direct modules/subpackages of the
package at the given path.
Arguments
---------
file_path : str
Path of the file to navigate the directory of. Typically the
`__init__.py` path this is called from, using `__file__`.
find_subpackages : bool
Whether we should find the subpackages as well.
Returns
-------
imports : List[str]
List of importable scripts with the same module.
"""
imports = []
module_dir = os.path.dirname(file_path)
for filename in os.listdir(module_dir):
if filename.startswith("__"):
continue
if filename.endswith(".py"):
imports.append(filename[:-3])
if find_subpackages and os.path.isdir(
os.path.join(module_dir, filename)
):
imports.append(filename)
return imports
def lazy_export(name: str, package: str):
"""Makes `name` lazily available under the module list for the specified
`package`, unless it was loaded already, in which case it is ignored.
Arguments
---------
name : str
Name of the module, as long as it can get imported with
`{package}.{name}`.
package : str
The relevant package, usually determined with `__name__` from the
`__init__.py`.
Returns
-------
None
"""
# already imported for real (e.g. utils.importutils itself)
if hasattr(sys.modules[package], name):
return
setattr(sys.modules[package], name, LazyModule(name, name, package))
def lazy_export_all(
init_file_path: str, package: str, export_subpackages: bool = False
):
"""Makes all modules under a module lazily importable merely by accessing
them; e.g. `foo/bar.py` could be accessed with `foo.bar.some_func()`.
Arguments
---------
init_file_path : str
Path of the `__init__.py` file, usually determined with `__file__` from
there.
package : str
The relevant package, usually determined with `__name__` from the
`__init__.py`.
export_subpackages : bool
Whether we should make the subpackages (subdirectories) available
directly as well.
"""
for name in find_imports(
init_file_path, find_subpackages=export_subpackages
):
lazy_export(name, package)
def deprecated_redirect(
old_import: str,
new_import: str,
extra_reason: Optional[str] = None,
also_lazy_export: bool = False,
) -> None:
"""Patches the module list to add a lazy redirection from `old_import` to
`new_import`, emitting a `DeprecationWarning` when imported.
Arguments
---------
old_import : str
Old module import path e.g. `mypackage.myoldmodule`
new_import : str
New module import path e.g. `mypackage.mycoolpackage.mynewmodule`
extra_reason : str, optional
If specified, extra text to attach to the warning for clarification
(e.g. justifying why the move has occurred, or additional problems to
look out for).
also_lazy_export : bool
Whether the module should also be exported as a lazy module in the
package determined in `old_import`.
e.g. if you had a `foo.bar.somefunc` import as `old_import`, assuming
you have `foo` imported (or lazy loaded), you could use
`foo.bar.somefunc` directly without importing `foo.bar` explicitly.
"""
redirect = DeprecatedModuleRedirect(
old_import, new_import, extra_reason=extra_reason
)
sys.modules[old_import] = redirect
if also_lazy_export:
package_sep_idx = old_import.rfind(".")
old_package = old_import[:package_sep_idx]
old_module = old_import[package_sep_idx + 1 :]
if not hasattr(sys.modules[old_package], old_module):
setattr(sys.modules[old_package], old_module, redirect)