_check_module.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. """
  2. Utility for locating a module (or package's __main__.py) with a given name
  3. and verifying it contains the PYTHON_ARGCOMPLETE_OK marker.
  4. The module name should be specified in a form usable with `python -m`.
  5. Intended to be invoked by argcomplete's global completion function.
  6. """
  7. import os
  8. import sys
  9. import tokenize
  10. from importlib.util import find_spec
  11. class ArgcompleteMarkerNotFound(RuntimeError):
  12. pass
  13. def find(name, return_package=False):
  14. names = name.split(".")
  15. # Look for the first importlib ModuleSpec that has `origin` set, indicating it's not a namespace package.
  16. for package_name_boundary in range(len(names)):
  17. spec = find_spec(".".join(names[: package_name_boundary + 1]))
  18. if spec is not None and spec.origin is not None:
  19. break
  20. if spec is None:
  21. raise ArgcompleteMarkerNotFound('no module named "{}"'.format(names[0]))
  22. if not spec.has_location:
  23. raise ArgcompleteMarkerNotFound("cannot locate file")
  24. if spec.submodule_search_locations is None:
  25. if len(names) != 1:
  26. raise ArgcompleteMarkerNotFound("{} is not a package".format(names[0]))
  27. return spec.origin
  28. if len(spec.submodule_search_locations) != 1:
  29. raise ArgcompleteMarkerNotFound("expecting one search location")
  30. path = os.path.join(spec.submodule_search_locations[0], *names[package_name_boundary + 1 :])
  31. if os.path.isdir(path):
  32. filename = "__main__.py"
  33. if return_package:
  34. filename = "__init__.py"
  35. return os.path.join(path, filename)
  36. else:
  37. return path + ".py"
  38. def main():
  39. try:
  40. name = sys.argv[1]
  41. except IndexError:
  42. raise ArgcompleteMarkerNotFound("missing argument on the command line")
  43. filename = find(name)
  44. try:
  45. fp = tokenize.open(filename)
  46. except OSError:
  47. raise ArgcompleteMarkerNotFound("cannot open file")
  48. with fp:
  49. head = fp.read(1024)
  50. if "PYTHON_ARGCOMPLETE_OK" not in head:
  51. raise ArgcompleteMarkerNotFound("marker not found")
  52. if __name__ == "__main__":
  53. try:
  54. main()
  55. except ArgcompleteMarkerNotFound as e:
  56. sys.exit(str(e))