_check_console_script.py 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. """
  2. Utility for locating the module (or package's __init__.py)
  3. associated with a given console_script name
  4. and verifying it contains the PYTHON_ARGCOMPLETE_OK marker.
  5. Such scripts are automatically generated and cannot contain
  6. the marker themselves, so we defer to the containing module or package.
  7. For more information on setuptools console_scripts, see
  8. https://setuptools.readthedocs.io/en/latest/setuptools.html#automatic-script-creation
  9. Intended to be invoked by argcomplete's global completion function.
  10. """
  11. import os
  12. import sys
  13. from importlib.metadata import EntryPoint
  14. from importlib.metadata import entry_points as importlib_entry_points
  15. from typing import Iterable
  16. from ._check_module import ArgcompleteMarkerNotFound, find
  17. def main():
  18. # Argument is the full path to the console script.
  19. script_path = sys.argv[1]
  20. # Find the module and function names that correspond to this
  21. # assuming it is actually a console script.
  22. name = os.path.basename(script_path)
  23. entry_points: Iterable[EntryPoint] = importlib_entry_points() # type:ignore
  24. # Python 3.12+ returns a tuple of entry point objects
  25. # whereas <=3.11 returns a SelectableGroups object
  26. if sys.version_info < (3, 12):
  27. entry_points = entry_points["console_scripts"] # type:ignore
  28. entry_points = [ep for ep in entry_points if ep.name == name and ep.group == "console_scripts"] # type:ignore
  29. if not entry_points:
  30. raise ArgcompleteMarkerNotFound("no entry point found matching script")
  31. entry_point = entry_points[0]
  32. module_name, function_name = entry_point.value.split(":", 1)
  33. # Check this looks like the script we really expected.
  34. with open(script_path) as f:
  35. script = f.read()
  36. if "from {} import {}".format(module_name, function_name) not in script:
  37. raise ArgcompleteMarkerNotFound("does not appear to be a console script")
  38. if "sys.exit({}())".format(function_name) not in script:
  39. raise ArgcompleteMarkerNotFound("does not appear to be a console script")
  40. # Look for the argcomplete marker in the script it imports.
  41. with open(find(module_name, return_package=True)) as f:
  42. head = f.read(1024)
  43. if "PYTHON_ARGCOMPLETE_OK" not in head:
  44. raise ArgcompleteMarkerNotFound("marker not found")
  45. if __name__ == "__main__":
  46. try:
  47. main()
  48. except ArgcompleteMarkerNotFound as e:
  49. sys.exit(str(e))