Add shared mapping key reader and reuse in tools

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

hyperbrowser/mapping_utils.py

@@ -1,18 +1,16 @@
 from collections.abc import Mapping as MappingABC
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, List
 
 from hyperbrowser.exceptions import HyperbrowserError
 
 
-def read_string_key_mapping(
+def read_string_mapping_keys(
     mapping_value: Any,
     *,
     expected_mapping_error: str,
     read_keys_error: str,
     non_string_key_error_builder: Callable[[object], str],
-    read_value_error_builder: Callable[[str], str],
-    key_display: Callable[[str], str],
-) -> Dict[str, object]:
+) -> List[str]:
     if not isinstance(mapping_value, MappingABC):
         raise HyperbrowserError(expected_mapping_error)
     try:
@@ -24,10 +22,30 @@ def read_string_key_mapping(
             read_keys_error,
             original_error=exc,
         ) from exc
+    normalized_keys: List[str] = []
     for key in mapping_keys:
         if type(key) is str:
+            normalized_keys.append(key)
             continue
         raise HyperbrowserError(non_string_key_error_builder(key))
+    return normalized_keys
+
+
+def read_string_key_mapping(
+    mapping_value: Any,
+    *,
+    expected_mapping_error: str,
+    read_keys_error: str,
+    non_string_key_error_builder: Callable[[object], str],
+    read_value_error_builder: Callable[[str], str],
+    key_display: Callable[[str], str],
+) -> Dict[str, object]:
+    mapping_keys = read_string_mapping_keys(
+        mapping_value,
+        expected_mapping_error=expected_mapping_error,
+        read_keys_error=read_keys_error,
+        non_string_key_error_builder=non_string_key_error_builder,
+    )
     normalized_mapping: Dict[str, object] = {}
     for key in mapping_keys:
         try:

hyperbrowser/tools/__init__.py

@@ -5,7 +5,10 @@
 
 from hyperbrowser.display_utils import format_string_key_for_error
 from hyperbrowser.exceptions import HyperbrowserError
-from hyperbrowser.mapping_utils import copy_mapping_values_by_string_keys
+from hyperbrowser.mapping_utils import (
+    copy_mapping_values_by_string_keys,
+    read_string_mapping_keys,
+)
 from hyperbrowser.models.agents.browser_use import StartBrowserUseTaskParams
 from hyperbrowser.models.crawl import StartCrawlJobParams
 from hyperbrowser.models.extract import StartExtractJobParams
@@ -60,20 +63,14 @@ def _format_tool_param_key_for_error(key: str) -> str:
 def _normalize_extract_schema_mapping(
     schema_value: MappingABC[object, Any],
 ) -> Dict[str, Any]:
-    try:
-        schema_keys = list(schema_value.keys())
-    except HyperbrowserError:
-        raise
-    except Exception as exc:
-        raise HyperbrowserError(
-            "Failed to read extract tool `schema` object keys",
-            original_error=exc,
-        ) from exc
-    normalized_schema_keys: list[str] = []
-    for key in schema_keys:
-        if type(key) is not str:
-            raise HyperbrowserError("Extract tool `schema` object keys must be strings")
-        normalized_schema_keys.append(key)
+    normalized_schema_keys = read_string_mapping_keys(
+        schema_value,
+        expected_mapping_error="Extract tool `schema` must be an object or JSON string",
+        read_keys_error="Failed to read extract tool `schema` object keys",
+        non_string_key_error_builder=lambda _key: (
+            "Extract tool `schema` object keys must be strings"
+        ),
+    )
     return copy_mapping_values_by_string_keys(
         schema_value,
         normalized_schema_keys,
@@ -114,67 +111,58 @@ def _prepare_extract_tool_params(params: Mapping[str, Any]) -> Dict[str, Any]:
 
 
 def _to_param_dict(params: Mapping[str, Any]) -> Dict[str, Any]:
-    if not isinstance(params, Mapping):
-        raise HyperbrowserError("tool params must be a mapping")
-    try:
-        param_keys = list(params.keys())
-    except HyperbrowserError:
-        raise
-    except Exception as exc:
-        raise HyperbrowserError(
-            "Failed to read tool params keys",
-            original_error=exc,
-        ) from exc
+    param_keys = read_string_mapping_keys(
+        params,
+        expected_mapping_error="tool params must be a mapping",
+        read_keys_error="Failed to read tool params keys",
+        non_string_key_error_builder=lambda _key: "tool params keys must be strings",
+    )
     for key in param_keys:
-        if type(key) is str:
-            try:
-                normalized_key = key.strip()
-                if type(normalized_key) is not str:
-                    raise TypeError("normalized tool param key must be a string")
-                is_empty_key = len(normalized_key) == 0
-            except HyperbrowserError:
-                raise
-            except Exception as exc:
-                raise HyperbrowserError(
-                    "Failed to normalize tool param key",
-                    original_error=exc,
-                ) from exc
-            if is_empty_key:
-                raise HyperbrowserError("tool params keys must not be empty")
-            try:
-                has_surrounding_whitespace = key != normalized_key
-            except HyperbrowserError:
-                raise
-            except Exception as exc:
-                raise HyperbrowserError(
-                    "Failed to normalize tool param key",
-                    original_error=exc,
-                ) from exc
-            if has_surrounding_whitespace:
-                raise HyperbrowserError(
-                    "tool params keys must not contain leading or trailing whitespace"
-                )
-            try:
-                contains_control_character = any(
-                    ord(character) < 32 or ord(character) == 127 for character in key
-                )
-            except HyperbrowserError:
-                raise
-            except Exception as exc:
-                raise HyperbrowserError(
-                    "Failed to validate tool param key characters",
-                    original_error=exc,
-                ) from exc
-            if contains_control_character:
-                raise HyperbrowserError(
-                    "tool params keys must not contain control characters"
-                )
-            continue
-        raise HyperbrowserError("tool params keys must be strings")
-    normalized_param_keys = [key for key in param_keys if type(key) is str]
+        try:
+            normalized_key = key.strip()
+            if type(normalized_key) is not str:
+                raise TypeError("normalized tool param key must be a string")
+            is_empty_key = len(normalized_key) == 0
+        except HyperbrowserError:
+            raise
+        except Exception as exc:
+            raise HyperbrowserError(
+                "Failed to normalize tool param key",
+                original_error=exc,
+            ) from exc
+        if is_empty_key:
+            raise HyperbrowserError("tool params keys must not be empty")
+        try:
+            has_surrounding_whitespace = key != normalized_key
+        except HyperbrowserError:
+            raise
+        except Exception as exc:
+            raise HyperbrowserError(
+                "Failed to normalize tool param key",
+                original_error=exc,
+            ) from exc
+        if has_surrounding_whitespace:
+            raise HyperbrowserError(
+                "tool params keys must not contain leading or trailing whitespace"
+            )
+        try:
+            contains_control_character = any(
+                ord(character) < 32 or ord(character) == 127 for character in key
+            )
+        except HyperbrowserError:
+            raise
+        except Exception as exc:
+            raise HyperbrowserError(
+                "Failed to validate tool param key characters",
+                original_error=exc,
+            ) from exc
+        if contains_control_character:
+            raise HyperbrowserError(
+                "tool params keys must not contain control characters"
+            )
     return copy_mapping_values_by_string_keys(
         params,
-        normalized_param_keys,
+        param_keys,
         read_value_error_builder=lambda key_display: (
             f"Failed to read tool param '{key_display}'"
         ),

tests/test_mapping_utils.py

@@ -5,6 +5,7 @@
 from hyperbrowser.exceptions import HyperbrowserError
 from hyperbrowser.mapping_utils import (
     copy_mapping_values_by_string_keys,
+    read_string_mapping_keys,
     read_string_key_mapping,
 )
 
@@ -116,6 +117,55 @@ def test_read_string_key_mapping_falls_back_for_unreadable_key_display():
         )
 
 
+def test_read_string_mapping_keys_returns_string_keys():
+    assert read_string_mapping_keys(
+        {"a": 1, "b": 2},
+        expected_mapping_error="expected mapping",
+        read_keys_error="failed keys",
+        non_string_key_error_builder=lambda key: (
+            f"non-string key: {type(key).__name__}"
+        ),
+    ) == ["a", "b"]
+
+
+def test_read_string_mapping_keys_rejects_non_mapping_values():
+    with pytest.raises(HyperbrowserError, match="expected mapping"):
+        read_string_mapping_keys(
+            ["a"],
+            expected_mapping_error="expected mapping",
+            read_keys_error="failed keys",
+            non_string_key_error_builder=lambda key: (
+                f"non-string key: {type(key).__name__}"
+            ),
+        )
+
+
+def test_read_string_mapping_keys_wraps_key_read_errors():
+    with pytest.raises(HyperbrowserError, match="failed keys") as exc_info:
+        read_string_mapping_keys(
+            _BrokenKeysMapping(),
+            expected_mapping_error="expected mapping",
+            read_keys_error="failed keys",
+            non_string_key_error_builder=lambda key: (
+                f"non-string key: {type(key).__name__}"
+            ),
+        )
+
+    assert isinstance(exc_info.value.original_error, RuntimeError)
+
+
+def test_read_string_mapping_keys_rejects_non_string_keys():
+    with pytest.raises(HyperbrowserError, match="non-string key: int"):
+        read_string_mapping_keys(
+            {1: "value"},
+            expected_mapping_error="expected mapping",
+            read_keys_error="failed keys",
+            non_string_key_error_builder=lambda key: (
+                f"non-string key: {type(key).__name__}"
+            ),
+        )
+
+
 def test_copy_mapping_values_by_string_keys_returns_selected_values():
     copied_values = copy_mapping_values_by_string_keys(
         {"field": "value", "other": "ignored"},