pyi_rth_pkgutil.py 3.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. #-----------------------------------------------------------------------------
  2. # Copyright (c) 2021-2023, PyInstaller Development Team.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. #
  7. # The full license is in the file COPYING.txt, distributed with this software.
  8. #
  9. # SPDX-License-Identifier: Apache-2.0
  10. #-----------------------------------------------------------------------------
  11. #
  12. # The run-time hook provides a custom module iteration function for our PyiFrozenFinder, which allows
  13. # `pkgutil.iter_modules()` to return entries for modules that are embedded in the PYZ archive. The non-embedded modules
  14. # (binary extensions, modules collected as only source .py files, etc.) are enumerated using the `fallback_finder`
  15. # provided by `PyiFrozenFinder` (which typically would be the python's `FileFinder`).
  16. def _pyi_rthook():
  17. import pkgutil
  18. import pyimod02_importers # PyInstaller's bootstrap module
  19. # This could, in fact, be implemented as `iter_modules()` method of the `PyiFrozenFinder`. However, we want to
  20. # avoid importing `pkgutil` in that bootstrap module (i.e., for the `pkgutil.iter_importer_modules()` call on the
  21. # fallback finder).
  22. def _iter_pyi_frozen_finder_modules(finder, prefix=''):
  23. # Fetch PYZ TOC tree from pyimod02_importers
  24. pyz_toc_tree = pyimod02_importers.get_pyz_toc_tree()
  25. # Finder has already pre-computed the package prefix implied by the search path. Use it to find the starting
  26. # node in the prefix tree.
  27. if finder._pyz_entry_prefix:
  28. pkg_name_parts = finder._pyz_entry_prefix.split('.')
  29. else:
  30. pkg_name_parts = []
  31. tree_node = pyz_toc_tree
  32. for pkg_name_part in pkg_name_parts:
  33. tree_node = tree_node.get(pkg_name_part)
  34. if not isinstance(tree_node, dict):
  35. # This check handles two cases:
  36. # a) path does not exist (`tree_node` is `None`)
  37. # b) path corresponds to a module instead of a package (`tree_node` is a leaf node (`str`)).
  38. tree_node = {}
  39. break
  40. # Dump the contents of the tree node.
  41. for entry_name, entry_data in tree_node.items():
  42. is_pkg = isinstance(entry_data, dict)
  43. yield prefix + entry_name, is_pkg
  44. # If our finder has a fall-back finder available, iterate its modules as well. By using the public
  45. # `fallback_finder` attribute, we force creation of the fallback finder as necessary.
  46. # NOTE: we do not care about potential duplicates here, because `pkgutil.iter_modules()` itself
  47. # keeps track of yielded names for purposes of de-duplication.
  48. if finder.fallback_finder is not None:
  49. yield from pkgutil.iter_importer_modules(finder.fallback_finder, prefix)
  50. pkgutil.iter_importer_modules.register(
  51. pyimod02_importers.PyiFrozenFinder,
  52. _iter_pyi_frozen_finder_modules,
  53. )
  54. _pyi_rthook()
  55. del _pyi_rthook