_helpers.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. #!/usr/bin/env python
  2. """Convenience functions and code used by ee/__init__.py.
  3. These functions are in general re-exported from the "ee" module and should be
  4. referenced from there (e.g. "ee.profilePrinting").
  5. """
  6. # Using lowercase function naming to match the JavaScript names.
  7. # pylint: disable=g-bad-name
  8. import contextlib
  9. import json
  10. import sys
  11. # pylint: disable=g-importing-member
  12. from . import data
  13. from . import oauth
  14. from .apifunction import ApiFunction
  15. # pylint: enable=g-importing-member
  16. from google.auth import crypt
  17. from google.oauth2 import service_account
  18. def ServiceAccountCredentials(email, key_file=None, key_data=None):
  19. """Configure OAuth2 credentials for a Google Service Account.
  20. Args:
  21. email: The email address of the account for which to configure credentials.
  22. Ignored if key_file or key_data represents a JSON service account key.
  23. key_file: The path to a file containing the private key associated with
  24. the service account. Both JSON and PEM files are supported.
  25. key_data: Raw key data to use, if key_file is not specified.
  26. Returns:
  27. An OAuth2 credentials object.
  28. """
  29. # Assume anything that doesn't end in '.pem' is a JSON key.
  30. if key_file and not key_file.endswith('.pem'):
  31. return service_account.Credentials.from_service_account_file(
  32. key_file, scopes=oauth.SCOPES)
  33. # If 'key_data' can be decoded as JSON, it's probably a raw JSON key.
  34. if key_data:
  35. try:
  36. key_data = json.loads(key_data)
  37. return service_account.Credentials.from_service_account_info(
  38. key_data, scopes=oauth.SCOPES)
  39. except ValueError:
  40. # It may actually be a raw PEM string, we'll try that below.
  41. pass
  42. # Probably a PEM key - just read the file into 'key_data'.
  43. if key_file:
  44. with open(key_file, 'r') as file_:
  45. key_data = file_.read()
  46. # Raw PEM key.
  47. signer = crypt.RSASigner.from_string(key_data)
  48. return service_account.Credentials(
  49. signer, email, oauth.TOKEN_URI, scopes=oauth.SCOPES)
  50. def call(func, *args, **kwargs):
  51. """Invoke the given algorithm with the specified args.
  52. Args:
  53. func: The function to call. Either an ee.Function object or the name of
  54. an API function.
  55. *args: The positional arguments to pass to the function.
  56. **kwargs: The named arguments to pass to the function.
  57. Returns:
  58. A ComputedObject representing the called function. If the signature
  59. specifies a recognized return type, the returned value will be cast
  60. to that type.
  61. """
  62. if isinstance(func, str):
  63. func = ApiFunction.lookup(func)
  64. return func.call(*args, **kwargs)
  65. def apply(func, named_args): # pylint: disable=redefined-builtin
  66. """Call a function with a dictionary of named arguments.
  67. Args:
  68. func: The function to call. Either an ee.Function object or the name of
  69. an API function.
  70. named_args: A dictionary of arguments to the function.
  71. Returns:
  72. A ComputedObject representing the called function. If the signature
  73. specifies a recognized return type, the returned value will be cast
  74. to that type.
  75. """
  76. if isinstance(func, str):
  77. func = ApiFunction.lookup(func)
  78. return func.apply(named_args)
  79. @contextlib.contextmanager
  80. def profilePrinting(destination=sys.stderr):
  81. # pylint: disable=g-doc-return-or-yield
  82. """Returns a context manager that prints a profile of enclosed API calls.
  83. The profile will be printed when the context ends, whether or not any error
  84. occurred within the context.
  85. # Simple example:
  86. with ee.profilePrinting():
  87. print ee.Number(1).add(1).getInfo()
  88. Args:
  89. destination: A file-like object to which the profile text is written.
  90. Defaults to sys.stderr.
  91. """
  92. # TODO(user): Figure out why ee.Profile.getProfiles isn't generated and fix
  93. # that.
  94. getProfiles = ApiFunction.lookup('Profile.getProfiles')
  95. profile_ids = []
  96. try:
  97. with data.profiling(profile_ids.append):
  98. yield
  99. finally:
  100. profile_text = getProfiles.call(ids=profile_ids).getInfo()
  101. destination.write(profile_text)