add counter

Author: jhamman

src/zarr/store/logging.py

@@ -3,10 +3,11 @@
 import inspect
 import logging
 import time
+from collections import defaultdict
 from contextlib import contextmanager
 from typing import TYPE_CHECKING
 
-from zarr.abc.store import Store
+from zarr.abc.store import AccessMode, Store
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Generator
@@ -16,6 +17,7 @@
 
 class LoggingStore(Store):
     _store: Store
+    counter: defaultdict[str, int]
 
     def __init__(
         self,
@@ -24,6 +26,7 @@ def __init__(
         log_handler: logging.Handler | None = None,
     ):
         self._store = store
+        self.counter = defaultdict(int)
 
         self._configure_logger(log_level, log_handler)
 
@@ -44,16 +47,19 @@ def _default_handler(self) -> logging.Handler:
         """Define a default log handler"""
         handler = logging.StreamHandler()
         handler.setLevel(self.log_level)
-        formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
-        handler.setFormatter(formatter)
+        handler.setFormatter(
+            logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        )
         return handler
 
     @contextmanager
     def log(self) -> Generator[None, None, None]:
-        op = f"{type(self._store).__name__}.{inspect.stack()[2].function}"
+        method = inspect.stack()[2].function
+        op = f"{type(self._store).__name__}.{method}"
         self.logger.info(f"Calling {op}")
         start_time = time.time()
         try:
+            self.counter[method] += 1
             yield
         finally:
             end_time = time.time()
@@ -79,6 +85,16 @@ def supports_listing(self) -> bool:
         with self.log():
             return self._store.supports_listing
 
+    @property
+    def _mode(self) -> AccessMode:  # type: ignore[override]
+        with self.log():
+            return self._store._mode
+
+    @property
+    def _is_open(self) -> bool:  # type: ignore[override]
+        with self.log():
+            return self._store._is_open
+
     async def empty(self) -> bool:
         with self.log():
             return await self._store.empty()

tests/v3/test_store/test_logging.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+import zarr
 from zarr.core.buffer import default_buffer_prototype
 from zarr.store.logging import LoggingStore
 
@@ -33,3 +34,17 @@ async def test_logging_store(store: Store, caplog) -> None:
         assert str(store) in tup[0]
     assert f"Calling {type(store).__name__}.list" in caplog.record_tuples[0][2]
     assert f"Finished {type(store).__name__}.list" in caplog.record_tuples[1][2]
+
+
+@pytest.mark.parametrize("store", ("local", "memory", "zip"), indirect=["store"])
+async def test_logging_store_counter(store: Store) -> None:
+    wrapped = LoggingStore(store=store, log_level="DEBUG")
+
+    arr = zarr.create(shape=(10,), store=wrapped, overwrite=True)
+    arr[:] = 1
+
+    assert wrapped.counter["set"] == 2
+    assert wrapped.counter["get"] == 0  # 1 if overwrite=False
+    assert wrapped.counter["list"] == 0
+    assert wrapped.counter["list_dir"] == 0
+    assert wrapped.counter["list_prefix"] == 0