Fix cooperation of our load(s) implementation with a user supplied object_hook

Author: Crimson-Crow

json_numpy.py

@@ -6,11 +6,14 @@
 import json
 from base64 import b64decode, b64encode
 from functools import partial
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 from numpy import frombuffer, generic, ndarray
 from numpy.lib.format import descr_to_dtype, dtype_to_descr
 
+if TYPE_CHECKING:  # pragma: no cover
+    from _typeshed import SupportsRead
+
 
 def default(
     o: Any, *, fallback_default: Callable[[Any], dict[str, Any]] | None = None
@@ -88,19 +91,25 @@ def dumps(*args: Any, cls: type[json.JSONEncoder] | None = None, **kwargs: Any)
     return _dumps(*args, cls=_patch_encoder, **kwargs)  # type: ignore[arg-type]
 
 
-def loads(*args: Any, **kwargs: Any) -> Any:
-    kwargs.setdefault("object_hook", object_hook)
-    return _loads(*args, **kwargs)
+def loads(
+    *args: Any, object_hook: Callable[[dict], Any] | None = None, **kwargs: Any
+) -> Any:
+    return _loads(
+        *args,
+        object_hook=_hook
+        if object_hook is None
+        else lambda dct: _hook(object_hook(dct)),
+        **kwargs,
+    )
 
 
 def dump(*args: Any, cls: type[json.JSONEncoder] | None = None, **kwargs: Any) -> None:
     kwargs["user_cls"] = cls
     return _dump(*args, cls=_patch_encoder, **kwargs)  # type: ignore[arg-type]
 
 
-def load(*args: Any, **kwargs: Any) -> Any:
-    kwargs.setdefault("object_hook", object_hook)
-    return _load(*args, **kwargs)
+def load(fp: SupportsRead[str | bytes], **kwargs: Any) -> Any:
+    return loads(fp.read(), **kwargs)
 
 
 def patch() -> None:

tests/test_json_numpy.py

@@ -221,3 +221,15 @@ def default(self, _: Any) -> dict[str, Any]:
         x = np.random.rand(5).astype(np.float32)
         dumped = json.dumps(x, cls=Encoder)
         self.assert_equal_with_type(json.loads(dumped), x)
+
+    def test_loads_object_hook(self) -> None:
+        def hook(dct: dict) -> dict | int:
+            if "foo" in dct:
+                return dct["foo"]
+            return dct
+
+        foo = {"foo": "bar"}
+        x = np.random.rand(5).astype(np.float32)
+        result = json.loads(json.dumps([foo, x]), object_hook=hook)
+        self.assertEqual(result[0], "bar")
+        self.assert_equal_with_type(result[1], x)