manager.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447
  1. # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
  2. # For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
  3. # Copyright (c) https://github.com/PyCQA/astroid/blob/main/CONTRIBUTORS.txt
  4. """astroid manager: avoid multiple astroid build of a same module when
  5. possible by providing a class responsible to get astroid representation
  6. from various source and using a cache of built modules)
  7. """
  8. from __future__ import annotations
  9. import collections
  10. import os
  11. import types
  12. import zipimport
  13. from collections.abc import Callable, Iterator, Sequence
  14. from importlib.util import find_spec, module_from_spec
  15. from typing import Any, ClassVar
  16. from astroid import nodes
  17. from astroid._cache import CACHE_MANAGER
  18. from astroid.const import BRAIN_MODULES_DIRECTORY
  19. from astroid.context import InferenceContext, _invalidate_cache
  20. from astroid.exceptions import AstroidBuildingError, AstroidImportError
  21. from astroid.interpreter._import import spec, util
  22. from astroid.modutils import (
  23. NoSourceFile,
  24. _cache_normalize_path_,
  25. file_info_from_modpath,
  26. get_source_file,
  27. is_module_name_part_of_extension_package_whitelist,
  28. is_python_source,
  29. is_stdlib_module,
  30. load_module_from_name,
  31. modpath_from_file,
  32. )
  33. from astroid.transforms import TransformVisitor
  34. from astroid.typing import AstroidManagerBrain, InferenceResult
  35. ZIP_IMPORT_EXTS = (".zip", ".egg", ".whl", ".pyz", ".pyzw")
  36. def safe_repr(obj: Any) -> str:
  37. try:
  38. return repr(obj)
  39. except Exception: # pylint: disable=broad-except
  40. return "???"
  41. class AstroidManager:
  42. """Responsible to build astroid from files or modules.
  43. Use the Borg (singleton) pattern.
  44. """
  45. name = "astroid loader"
  46. brain: AstroidManagerBrain = {
  47. "astroid_cache": {},
  48. "_mod_file_cache": {},
  49. "_failed_import_hooks": [],
  50. "always_load_extensions": False,
  51. "optimize_ast": False,
  52. "extension_package_whitelist": set(),
  53. "_transform": TransformVisitor(),
  54. }
  55. max_inferable_values: ClassVar[int] = 100
  56. def __init__(self) -> None:
  57. # NOTE: cache entries are added by the [re]builder
  58. self.astroid_cache = AstroidManager.brain["astroid_cache"]
  59. self._mod_file_cache = AstroidManager.brain["_mod_file_cache"]
  60. self._failed_import_hooks = AstroidManager.brain["_failed_import_hooks"]
  61. self.always_load_extensions = AstroidManager.brain["always_load_extensions"]
  62. self.optimize_ast = AstroidManager.brain["optimize_ast"]
  63. self.extension_package_whitelist = AstroidManager.brain[
  64. "extension_package_whitelist"
  65. ]
  66. self._transform = AstroidManager.brain["_transform"]
  67. @property
  68. def register_transform(self):
  69. # This and unregister_transform below are exported for convenience
  70. return self._transform.register_transform
  71. @property
  72. def unregister_transform(self):
  73. return self._transform.unregister_transform
  74. @property
  75. def builtins_module(self) -> nodes.Module:
  76. return self.astroid_cache["builtins"]
  77. def visit_transforms(self, node: nodes.NodeNG) -> InferenceResult:
  78. """Visit the transforms and apply them to the given *node*."""
  79. return self._transform.visit(node)
  80. def ast_from_file(
  81. self,
  82. filepath: str,
  83. modname: str | None = None,
  84. fallback: bool = True,
  85. source: bool = False,
  86. ) -> nodes.Module:
  87. """Given a module name, return the astroid object."""
  88. try:
  89. filepath = get_source_file(filepath, include_no_ext=True)
  90. source = True
  91. except NoSourceFile:
  92. pass
  93. if modname is None:
  94. try:
  95. modname = ".".join(modpath_from_file(filepath))
  96. except ImportError:
  97. modname = filepath
  98. if (
  99. modname in self.astroid_cache
  100. and self.astroid_cache[modname].file == filepath
  101. ):
  102. return self.astroid_cache[modname]
  103. if source:
  104. # pylint: disable=import-outside-toplevel; circular import
  105. from astroid.builder import AstroidBuilder
  106. return AstroidBuilder(self).file_build(filepath, modname)
  107. if fallback and modname:
  108. return self.ast_from_module_name(modname)
  109. raise AstroidBuildingError("Unable to build an AST for {path}.", path=filepath)
  110. def ast_from_string(
  111. self, data: str, modname: str = "", filepath: str | None = None
  112. ) -> nodes.Module:
  113. """Given some source code as a string, return its corresponding astroid
  114. object.
  115. """
  116. # pylint: disable=import-outside-toplevel; circular import
  117. from astroid.builder import AstroidBuilder
  118. return AstroidBuilder(self).string_build(data, modname, filepath)
  119. def _build_stub_module(self, modname: str) -> nodes.Module:
  120. # pylint: disable=import-outside-toplevel; circular import
  121. from astroid.builder import AstroidBuilder
  122. return AstroidBuilder(self).string_build("", modname)
  123. def _build_namespace_module(
  124. self, modname: str, path: Sequence[str]
  125. ) -> nodes.Module:
  126. # pylint: disable=import-outside-toplevel; circular import
  127. from astroid.builder import build_namespace_package_module
  128. return build_namespace_package_module(modname, path)
  129. def _can_load_extension(self, modname: str) -> bool:
  130. if self.always_load_extensions:
  131. return True
  132. if is_stdlib_module(modname):
  133. return True
  134. return is_module_name_part_of_extension_package_whitelist(
  135. modname, self.extension_package_whitelist
  136. )
  137. def ast_from_module_name( # noqa: C901
  138. self,
  139. modname: str | None,
  140. context_file: str | None = None,
  141. use_cache: bool = True,
  142. ) -> nodes.Module:
  143. """Given a module name, return the astroid object."""
  144. if modname is None:
  145. raise AstroidBuildingError("No module name given.")
  146. # Sometimes we don't want to use the cache. For example, when we're
  147. # importing a module with the same name as the file that is importing
  148. # we want to fallback on the import system to make sure we get the correct
  149. # module.
  150. if modname in self.astroid_cache and use_cache:
  151. return self.astroid_cache[modname]
  152. if modname == "__main__":
  153. return self._build_stub_module(modname)
  154. if context_file:
  155. old_cwd = os.getcwd()
  156. os.chdir(os.path.dirname(context_file))
  157. try:
  158. found_spec = self.file_from_module_name(modname, context_file)
  159. if found_spec.type == spec.ModuleType.PY_ZIPMODULE:
  160. module = self.zip_import_data(found_spec.location)
  161. if module is not None:
  162. return module
  163. elif found_spec.type in (
  164. spec.ModuleType.C_BUILTIN,
  165. spec.ModuleType.C_EXTENSION,
  166. ):
  167. if (
  168. found_spec.type == spec.ModuleType.C_EXTENSION
  169. and not self._can_load_extension(modname)
  170. ):
  171. return self._build_stub_module(modname)
  172. try:
  173. named_module = load_module_from_name(modname)
  174. except Exception as e:
  175. raise AstroidImportError(
  176. "Loading {modname} failed with:\n{error}",
  177. modname=modname,
  178. path=found_spec.location,
  179. ) from e
  180. return self.ast_from_module(named_module, modname)
  181. elif found_spec.type == spec.ModuleType.PY_COMPILED:
  182. raise AstroidImportError(
  183. "Unable to load compiled module {modname}.",
  184. modname=modname,
  185. path=found_spec.location,
  186. )
  187. elif found_spec.type == spec.ModuleType.PY_NAMESPACE:
  188. return self._build_namespace_module(
  189. modname, found_spec.submodule_search_locations or []
  190. )
  191. elif found_spec.type == spec.ModuleType.PY_FROZEN:
  192. if found_spec.location is None:
  193. return self._build_stub_module(modname)
  194. # For stdlib frozen modules we can determine the location and
  195. # can therefore create a module from the source file
  196. return self.ast_from_file(found_spec.location, modname, fallback=False)
  197. if found_spec.location is None:
  198. raise AstroidImportError(
  199. "Can't find a file for module {modname}.", modname=modname
  200. )
  201. return self.ast_from_file(found_spec.location, modname, fallback=False)
  202. except AstroidBuildingError as e:
  203. for hook in self._failed_import_hooks:
  204. try:
  205. return hook(modname)
  206. except AstroidBuildingError:
  207. pass
  208. raise e
  209. finally:
  210. if context_file:
  211. os.chdir(old_cwd)
  212. def zip_import_data(self, filepath: str) -> nodes.Module | None:
  213. if zipimport is None:
  214. return None
  215. # pylint: disable=import-outside-toplevel; circular import
  216. from astroid.builder import AstroidBuilder
  217. builder = AstroidBuilder(self)
  218. for ext in ZIP_IMPORT_EXTS:
  219. try:
  220. eggpath, resource = filepath.rsplit(ext + os.path.sep, 1)
  221. except ValueError:
  222. continue
  223. try:
  224. # pylint: disable-next=no-member
  225. importer = zipimport.zipimporter(eggpath + ext)
  226. zmodname = resource.replace(os.path.sep, ".")
  227. if importer.is_package(resource):
  228. zmodname = zmodname + ".__init__"
  229. module = builder.string_build(
  230. importer.get_source(resource), zmodname, filepath
  231. )
  232. return module
  233. except Exception: # pylint: disable=broad-except
  234. continue
  235. return None
  236. def file_from_module_name(
  237. self, modname: str, contextfile: str | None
  238. ) -> spec.ModuleSpec:
  239. try:
  240. value = self._mod_file_cache[(modname, contextfile)]
  241. except KeyError:
  242. try:
  243. value = file_info_from_modpath(
  244. modname.split("."), context_file=contextfile
  245. )
  246. except ImportError as e:
  247. # pylint: disable-next=redefined-variable-type
  248. value = AstroidImportError(
  249. "Failed to import module {modname} with error:\n{error}.",
  250. modname=modname,
  251. # we remove the traceback here to save on memory usage (since these exceptions are cached)
  252. error=e.with_traceback(None),
  253. )
  254. self._mod_file_cache[(modname, contextfile)] = value
  255. if isinstance(value, AstroidBuildingError):
  256. # we remove the traceback here to save on memory usage (since these exceptions are cached)
  257. raise value.with_traceback(None) # pylint: disable=no-member
  258. return value
  259. def ast_from_module(
  260. self, module: types.ModuleType, modname: str | None = None
  261. ) -> nodes.Module:
  262. """Given an imported module, return the astroid object."""
  263. modname = modname or module.__name__
  264. if modname in self.astroid_cache:
  265. return self.astroid_cache[modname]
  266. try:
  267. # some builtin modules don't have __file__ attribute
  268. filepath = module.__file__
  269. if is_python_source(filepath):
  270. # Type is checked in is_python_source
  271. return self.ast_from_file(filepath, modname) # type: ignore[arg-type]
  272. except AttributeError:
  273. pass
  274. # pylint: disable=import-outside-toplevel; circular import
  275. from astroid.builder import AstroidBuilder
  276. return AstroidBuilder(self).module_build(module, modname)
  277. def ast_from_class(self, klass: type, modname: str | None = None) -> nodes.ClassDef:
  278. """Get astroid for the given class."""
  279. if modname is None:
  280. try:
  281. modname = klass.__module__
  282. except AttributeError as exc:
  283. raise AstroidBuildingError(
  284. "Unable to get module for class {class_name}.",
  285. cls=klass,
  286. class_repr=safe_repr(klass),
  287. modname=modname,
  288. ) from exc
  289. modastroid = self.ast_from_module_name(modname)
  290. ret = modastroid.getattr(klass.__name__)[0]
  291. assert isinstance(ret, nodes.ClassDef)
  292. return ret
  293. def infer_ast_from_something(
  294. self, obj: object, context: InferenceContext | None = None
  295. ) -> Iterator[InferenceResult]:
  296. """Infer astroid for the given class."""
  297. if hasattr(obj, "__class__") and not isinstance(obj, type):
  298. klass = obj.__class__
  299. elif isinstance(obj, type):
  300. klass = obj
  301. else:
  302. raise AstroidBuildingError( # pragma: no cover
  303. "Unable to get type for {class_repr}.",
  304. cls=None,
  305. class_repr=safe_repr(obj),
  306. )
  307. try:
  308. modname = klass.__module__
  309. except AttributeError as exc:
  310. raise AstroidBuildingError(
  311. "Unable to get module for {class_repr}.",
  312. cls=klass,
  313. class_repr=safe_repr(klass),
  314. ) from exc
  315. except Exception as exc:
  316. raise AstroidImportError(
  317. "Unexpected error while retrieving module for {class_repr}:\n"
  318. "{error}",
  319. cls=klass,
  320. class_repr=safe_repr(klass),
  321. ) from exc
  322. try:
  323. name = klass.__name__
  324. except AttributeError as exc:
  325. raise AstroidBuildingError(
  326. "Unable to get name for {class_repr}:\n",
  327. cls=klass,
  328. class_repr=safe_repr(klass),
  329. ) from exc
  330. except Exception as exc:
  331. raise AstroidImportError(
  332. "Unexpected error while retrieving name for {class_repr}:\n{error}",
  333. cls=klass,
  334. class_repr=safe_repr(klass),
  335. ) from exc
  336. # take care, on living object __module__ is regularly wrong :(
  337. modastroid = self.ast_from_module_name(modname)
  338. if klass is obj:
  339. for inferred in modastroid.igetattr(name, context):
  340. yield inferred
  341. else:
  342. for inferred in modastroid.igetattr(name, context):
  343. yield inferred.instantiate_class()
  344. def register_failed_import_hook(self, hook: Callable[[str], nodes.Module]) -> None:
  345. """Registers a hook to resolve imports that cannot be found otherwise.
  346. `hook` must be a function that accepts a single argument `modname` which
  347. contains the name of the module or package that could not be imported.
  348. If `hook` can resolve the import, must return a node of type `astroid.Module`,
  349. otherwise, it must raise `AstroidBuildingError`.
  350. """
  351. self._failed_import_hooks.append(hook)
  352. def cache_module(self, module: nodes.Module) -> None:
  353. """Cache a module if no module with the same name is known yet."""
  354. self.astroid_cache.setdefault(module.name, module)
  355. def bootstrap(self) -> None:
  356. """Bootstrap the required AST modules needed for the manager to work.
  357. The bootstrap usually involves building the AST for the builtins
  358. module, which is required by the rest of astroid to work correctly.
  359. """
  360. from astroid import raw_building # pylint: disable=import-outside-toplevel
  361. raw_building._astroid_bootstrapping()
  362. def clear_cache(self) -> None:
  363. """Clear the underlying caches, bootstrap the builtins module and
  364. re-register transforms.
  365. """
  366. # import here because of cyclic imports
  367. # pylint: disable=import-outside-toplevel
  368. from astroid.inference_tip import clear_inference_tip_cache
  369. from astroid.interpreter.objectmodel import ObjectModel
  370. from astroid.nodes.node_classes import LookupMixIn
  371. from astroid.nodes.scoped_nodes import ClassDef
  372. clear_inference_tip_cache()
  373. _invalidate_cache() # inference context cache
  374. self.astroid_cache.clear()
  375. # NB: not a new TransformVisitor()
  376. AstroidManager.brain["_transform"].transforms = collections.defaultdict(list)
  377. CACHE_MANAGER.clear_all_caches()
  378. for lru_cache in (
  379. LookupMixIn.lookup,
  380. _cache_normalize_path_,
  381. util.is_namespace,
  382. ObjectModel.attributes,
  383. ClassDef._metaclass_lookup_attribute,
  384. ):
  385. lru_cache.cache_clear() # type: ignore[attr-defined]
  386. self.bootstrap()
  387. # Reload brain plugins. During initialisation this is done in astroid.__init__.py
  388. for module in BRAIN_MODULES_DIRECTORY.iterdir():
  389. if module.suffix == ".py":
  390. module_spec = find_spec(f"astroid.brain.{module.stem}")
  391. assert module_spec
  392. module_object = module_from_spec(module_spec)
  393. assert module_spec.loader
  394. module_spec.loader.exec_module(module_object)