From ca0e130091d1e99ed8ee36395edb91c677d1727a Mon Sep 17 00:00:00 2001
From: Yiling-J <njjyl723@gmail.com>
Date: Sun, 12 Nov 2023 10:03:00 +0800
Subject: [PATCH] add cache stats api (#20)

---
 README.md           | 13 ++++++++++++-
 tests/test_cache.py | 21 +++++++++++++++++++++
 theine/models.py    |  8 ++++++--
 theine/theine.py    |  9 +++++++++
 4 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index de37c87..6029083 100644
--- a/README.md
+++ b/README.md
@@ -95,6 +95,17 @@ cache.close()
 
 # clear cache
 cache.clear()
+
+# get current cache stats, please call stats() again if you need updated stats
+stats = cache.stats()
+print(stats.request_count, stats.hit_count, stats.hit_rate)
+
+# get cache max size
+cache.max_size
+
+# get cache current size
+len(cache)
+
 ```
 
 ## Decorator
@@ -240,6 +251,6 @@ Meta shared anonymized trace captured from large scale production cache services
 ![hit ratios](benchmarks/fb.png)
 
 ## Support
-Open an issue, ask question in discussions or join discord channel: https://discord.gg/StrgfPaQqE 
+Open an issue, ask question in discussions or join discord channel: https://discord.gg/StrgfPaQqE
 
 Theine Go version is also available, which focus on concurrency performance, take a look if you are interested: [Theine Go](https://github.com/Yiling-J/theine-go).
diff --git a/tests/test_cache.py b/tests/test_cache.py
index 053bee0..058af2e 100644
--- a/tests/test_cache.py
+++ b/tests/test_cache.py
@@ -1,6 +1,7 @@
 from datetime import timedelta
 from random import randint
 from time import sleep
+from bounded_zipf import Zipf
 
 import pytest
 
@@ -137,3 +138,23 @@ def test_close_cache(policy):
         cache.set("foo", "bar", timedelta(seconds=60))
         cache.close()
         assert cache._maintainer.is_alive() is False
+
+
+def test_cache_stats(policy):
+    cache = Cache(policy, 5000)
+    assert cache.max_size == 5000
+    assert len(cache) == 0
+    z = Zipf(1.0001, 10, 20000)
+    for _ in range(20000):
+        i = z.get()
+        key = f"key:{i}"
+        v = cache.get(key)
+        if v is None:
+            cache.set(key, key)
+    stats = cache.stats()
+    assert stats.hit_count > 0
+    assert stats.miss_count > 0
+    assert stats.request_count == stats.hit_count + stats.miss_count
+    assert stats.hit_rate > 0.5
+    assert stats.hit_rate < 1
+    assert stats.hit_rate == stats.hit_count / stats.request_count
diff --git a/theine/models.py b/theine/models.py
index 9f51cfe..07088d0 100644
--- a/theine/models.py
+++ b/theine/models.py
@@ -1,2 +1,6 @@
-from dataclasses import dataclass
-from typing import Any, Optional
+class CacheStats:
+    def __init__(self, total: int, hit: int):
+        self.request_count = total
+        self.hit_count = hit
+        self.miss_count = self.request_count - self.hit_count
+        self.hit_rate = self.hit_count / self.request_count
diff --git a/theine/theine.py b/theine/theine.py
index eec2744..9d69369 100644
--- a/theine/theine.py
+++ b/theine/theine.py
@@ -24,6 +24,7 @@
 from typing_extensions import ParamSpec, Protocol
 
 from theine.exceptions import InvalidTTL
+from theine.models import CacheStats
 
 sentinel = object()
 
@@ -277,6 +278,9 @@ def __init__(self, policy: str, size: int):
         self._closed = False
         self._maintainer = Thread(target=self.maintenance, daemon=True)
         self._maintainer.start()
+        self._total = 0
+        self._hit = 0
+        self.max_size = size
 
     def __len__(self) -> int:
         return self.core.len()
@@ -288,6 +292,7 @@ def get(self, key: Hashable, default: Any = None) -> Any:
         :param key: key hashable, use str/int for best performance.
         :param default: returned value if key is not found in cache, default None.
         """
+        self._total += 1
         auto_key = False
         key_str = ""
         if isinstance(key, str):
@@ -304,6 +309,7 @@ def get(self, key: Hashable, default: Any = None) -> Any:
                 self.key_gen.remove(key_str)
             return default
 
+        self._hit += 1
         return self._cache[index]
 
     def _access(self, key: Hashable, ttl: Optional[timedelta] = None):
@@ -438,3 +444,6 @@ def close(self):
     def __del__(self):
         self.clear()
         self.close()
+
+    def stats(self) -> CacheStats:
+        return CacheStats(self._total, self._hit)