Rename get_api_key to get_environment_variable

Renames the `get_api_key` function in `coreutils.py` to `get_environment_variable` to make it more generic and applicable to a wider range of use cases.

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."

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:

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.")

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")