From a5e7c33517c5b005d490fed099cb44987ff76acd Mon Sep 17 00:00:00 2001 From: Xee authors Date: Sun, 29 Dec 2024 01:15:20 -0800 Subject: [PATCH] Add argument to generate ee credentials on-demand. PiperOrigin-RevId: 710437577 --- xee/ext.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/xee/ext.py b/xee/ext.py index 763fe84..6b83da1 100644 --- a/xee/ext.py +++ b/xee/ext.py @@ -19,6 +19,7 @@ from __future__ import annotations import concurrent.futures +import copy import functools import importlib import itertools @@ -26,7 +27,7 @@ import math import os import sys -from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union from urllib import parse import warnings @@ -803,7 +804,20 @@ def _ee_init_check(self): 'Attempting to initialize using application default credentials.' ) - ee.Initialize(**(self.store.ee_init_kwargs or {})) + ee_init_kwargs = copy.copy(self.store.ee_init_kwargs) or {} + if ( + 'credentials' in ee_init_kwargs + and 'credentials_function' in ee_init_kwargs + ): + raise ValueError( + 'Cannot specify both credentials and credentials_function.' + ) + if 'credentials_function' in ee_init_kwargs: + credentials_function: Callable[[], Any] = ee_init_kwargs.pop( + 'credentials_function' + ) + ee_init_kwargs['credentials'] = credentials_function() + ee.Initialize(**ee_init_kwargs) def __getitem__(self, key: indexing.ExplicitIndexer) -> np.typing.ArrayLike: return indexing.explicit_indexing_adapter(