pyinfo.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. from __future__ import annotations
  2. """Utilities to find the site and prefix information of a Python executable.
  3. This file MUST remain compatible with all Python 3.8+ versions. Since we cannot make any
  4. assumptions about the Python being executed, this module should not use *any* dependencies outside
  5. of the standard library found in Python 3.8. This file is run each mypy run, so it should be kept
  6. as fast as possible.
  7. """
  8. import sys
  9. if __name__ == "__main__":
  10. # HACK: We don't want to pick up mypy.types as the top-level types
  11. # module. This could happen if this file is run as a script.
  12. # This workaround fixes this for Python versions before 3.11.
  13. if sys.version_info < (3, 11):
  14. old_sys_path = sys.path
  15. sys.path = sys.path[1:]
  16. import types # noqa: F401
  17. sys.path = old_sys_path
  18. import os
  19. import site
  20. import sysconfig
  21. def getsitepackages() -> list[str]:
  22. res = []
  23. if hasattr(site, "getsitepackages"):
  24. res.extend(site.getsitepackages())
  25. if hasattr(site, "getusersitepackages") and site.ENABLE_USER_SITE:
  26. res.insert(0, site.getusersitepackages())
  27. else:
  28. res = [sysconfig.get_paths()["purelib"]]
  29. return res
  30. def getsyspath() -> list[str]:
  31. # Do not include things from the standard library
  32. # because those should come from typeshed.
  33. stdlib_zip = os.path.join(
  34. sys.base_exec_prefix,
  35. getattr(sys, "platlibdir", "lib"),
  36. f"python{sys.version_info.major}{sys.version_info.minor}.zip",
  37. )
  38. stdlib = sysconfig.get_path("stdlib")
  39. stdlib_ext = os.path.join(stdlib, "lib-dynload")
  40. excludes = {stdlib_zip, stdlib, stdlib_ext}
  41. # Drop the first entry of sys.path
  42. # - If pyinfo.py is executed as a script (in a subprocess), this is the directory
  43. # containing pyinfo.py
  44. # - Otherwise, if mypy launched via console script, this is the directory of the script
  45. # - Otherwise, if mypy launched via python -m mypy, this is the current directory
  46. # In all these cases, it is desirable to drop the first entry
  47. # Note that mypy adds the cwd to SearchPaths.python_path, so we still find things on the
  48. # cwd consistently (the return value here sets SearchPaths.package_path)
  49. # Python 3.11 adds a "safe_path" flag wherein Python won't automatically prepend
  50. # anything to sys.path. In this case, the first entry of sys.path is no longer special.
  51. offset = 0 if sys.version_info >= (3, 11) and sys.flags.safe_path else 1
  52. abs_sys_path = (os.path.abspath(p) for p in sys.path[offset:])
  53. return [p for p in abs_sys_path if p not in excludes]
  54. def getsearchdirs() -> tuple[list[str], list[str]]:
  55. return (getsyspath(), getsitepackages())
  56. if __name__ == "__main__":
  57. if sys.argv[-1] == "getsearchdirs":
  58. print(repr(getsearchdirs()))
  59. else:
  60. print("ERROR: incorrect argument to pyinfo.py.", file=sys.stderr)
  61. sys.exit(1)