diff --git a/geemap/ai.py b/geemap/ai.py index ad2abeff3d..c45dde87ca 100644 --- a/geemap/ai.py +++ b/geemap/ai.py @@ -12,7 +12,7 @@ import ipywidgets as widgets from IPython.display import display, HTML from typing import Optional -from .common import get_api_key, temp_file_path +from .common import get_environment_variable, temp_file_path from .geemap import Map, ee_initialize try: @@ -57,14 +57,16 @@ def __init__( # Initialization if project is None: - project = get_api_key("EE_PROJECT_ID") or get_api_key("GOOGLE_PROJECT_ID") + project = get_environment_variable( + "EE_PROJECT_ID" + ) or get_environment_variable("GOOGLE_PROJECT_ID") if project is None: raise ValueError( "Please provide a valid project ID via the 'project' parameter." ) if google_api_key is None: - google_api_key = get_api_key("GOOGLE_API_KEY") + google_api_key = get_environment_variable("GOOGLE_API_KEY") if google_api_key is None: raise ValueError( "Please provide a valid Google API key via the 'google_api_key' parameter." diff --git a/geemap/coreutils.py b/geemap/coreutils.py index 4fde8b586d..0eea7db171 100644 --- a/geemap/coreutils.py +++ b/geemap/coreutils.py @@ -1,5 +1,6 @@ import json import os +import sys import zipfile import ee @@ -13,32 +14,29 @@ pass -def get_api_key(name: Optional[str] = None, key: Optional[str] = None) -> Optional[str]: - """ - Retrieves an API key. If a key is provided, it is returned directly. If a - name is provided, the function attempts to retrieve the key from user data - (if running in Google Colab) or from environment variables. +def get_environment_variable(key: str) -> Optional[str]: + """Retrieves an environment variable or Colab secret for the given key. + + Colab secrets have precedence over environment variables. Args: - name (Optional[str], optional): The name of the key to retrieve. Defaults to None. - key (Optional[str], optional): The key to return directly. Defaults to None. + key (str): The key that's used to fetch the environment variable. Returns: - Optional[str]: The retrieved key, or None if no key was found. + Optional[str]: The retrieved key, or None if no environment variable was found. """ + if not key: + return None - if key is not None: - return key - elif name is not None: - if in_colab_shell(): - from google.colab import userdata + if in_colab_shell(): + from google.colab import userdata - try: - return userdata.get(name) - except (userdata.SecretNotFoundError, userdata.NotebookAccessError): - return os.environ.get(name) - else: - return os.environ.get(name) + try: + return userdata.get(key) + except (userdata.SecretNotFoundError, userdata.NotebookAccessError): + pass + + return os.environ.get(key) def ee_initialize( @@ -85,7 +83,7 @@ def ee_initialize( kwargs["http_transport"] = httplib2.Http() if project is None: - kwargs["project"] = get_api_key("EE_PROJECT_ID") + kwargs["project"] = get_environment_variable("EE_PROJECT_ID") else: kwargs["project"] = project @@ -100,7 +98,7 @@ def ee_initialize( auth_args["auth_mode"] = auth_mode if ee.data._credentials is None: - ee_token = get_api_key(token_name) + ee_token = get_environment_variable(token_name) if service_account: try: credential_file_path = os.path.expanduser( @@ -342,7 +340,7 @@ def get_google_maps_api_key(key: str = "GOOGLE_MAPS_API_KEY") -> Optional[str]: Returns: str: The API key, or None if it could not be found. """ - if api_key := get_api_key(key): + if api_key := get_environment_variable(key): return api_key return os.environ.get(key, None) @@ -354,12 +352,7 @@ def in_colab_shell() -> bool: Returns: bool: True if running in Google Colab, False otherwise. """ - import sys - - if "google.colab" in sys.modules: - return True - else: - return False + return "google.colab" in sys.modules def check_color(in_color: Union[str, Tuple[int, int, int]]) -> str: diff --git a/geemap/maplibregl.py b/geemap/maplibregl.py index 3f7edccf60..88548b8a8b 100644 --- a/geemap/maplibregl.py +++ b/geemap/maplibregl.py @@ -1292,8 +1292,8 @@ def to_html( html = html.replace(div_before, div_after) if replace_key or (os.getenv("MAPTILER_REPLACE_KEY") is not None): - key_before = get_api_key("MAPTILER_KEY") - key_after = get_api_key("MAPTILER_KEY_PUBLIC") + key_before = get_environment_variable("MAPTILER_KEY") + key_after = get_environment_variable("MAPTILER_KEY_PUBLIC") if key_after is not None: html = html.replace(key_before, key_after) @@ -2217,7 +2217,7 @@ def _get_3d_terrain_style( """ if api_key is None: - api_key = get_api_key(token) + api_key = get_environment_variable(token) if api_key is None: print("An API key is required to use the 3D terrain feature.") @@ -2748,7 +2748,7 @@ def add_3d_buildings( None """ - MAPTILER_KEY = get_api_key("MAPTILER_KEY") + MAPTILER_KEY = get_environment_variable("MAPTILER_KEY") source = { "url": f"https://api.maptiler.com/tiles/v3/tiles.json?key={MAPTILER_KEY}", "type": "vector", @@ -2913,7 +2913,7 @@ def construct_maptiler_style(style: str, api_key: Optional[str] = None) -> str: """ if api_key is None: - api_key = get_api_key("MAPTILER_KEY") + api_key = get_environment_variable("MAPTILER_KEY") url = f"https://api.maptiler.com/maps/{style}/style.json?key={api_key}" @@ -2965,7 +2965,7 @@ def maptiler_3d_style( """ if api_key is None: - api_key = get_api_key(token) + api_key = get_environment_variable(token) if api_key is None: print("An API key is required to use the 3D terrain feature.") diff --git a/tests/test_coreutils.py b/tests/test_coreutils.py new file mode 100644 index 0000000000..a3c7ef29b5 --- /dev/null +++ b/tests/test_coreutils.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python + +"""Tests for `coreutils` module.""" +import os +import sys +import unittest +from unittest import mock + +from geemap import coreutils + + +class FakeSecretNotFoundError(Exception): + """google.colab.userdata.SecretNotFoundError fake.""" + + +class FakeNotebookAccessError(Exception): + """google.colab.userdata.NotebookAccessError fake.""" + + +class TestCoreUtils(unittest.TestCase): + """Tests for core utilss.""" + + def test_get_environment_invalid_key(self): + """Verifies None is returned if keys are invalid.""" + self.assertIsNone(coreutils.get_environment_variable(None)) + self.assertIsNone(coreutils.get_environment_variable("")) + + @mock.patch.dict(os.environ, {"key": "value"}) + def test_get_environment_variable_unknown_key(self): + """Verifies None is returned if the environment variable could not be found.""" + self.assertIsNone(coreutils.get_environment_variable("unknown-key")) + + @mock.patch.dict(os.environ, {"key": "value"}) + def test_get_environment_variable_from_env(self): + """Verifies environment variables are read from environment variables.""" + self.assertEqual(coreutils.get_environment_variable("key"), "value") + + @mock.patch.dict("sys.modules", {"google.colab": mock.Mock()}) + def test_get_environment_variable_from_colab(self): + """Verifies environment variables are read from Colab secrets.""" + mock_colab = sys.modules["google.colab"] + mock_colab.userdata.get.return_value = "colab-value" + + self.assertEqual(coreutils.get_environment_variable("key"), "colab-value") + mock_colab.userdata.get.assert_called_once_with("key") + + @mock.patch.dict(os.environ, {"key": "environ-value"}) + @mock.patch.dict("sys.modules", {"google.colab": mock.Mock()}) + def test_get_environment_variable_colab_fails_fallback_to_env(self): + """Verifies environment variables are read if a Colab secret read fails.""" + mock_colab = sys.modules["google.colab"] + mock_colab.userdata.SecretNotFoundError = FakeSecretNotFoundError + mock_colab.userdata.NotebookAccessError = FakeNotebookAccessError + mock_colab.userdata.get.side_effect = FakeNotebookAccessError() + + self.assertEqual(coreutils.get_environment_variable("key"), "environ-value")