From 40ab4cb388c33a56523535e19e2af4bfb8ea6f6d Mon Sep 17 00:00:00 2001
From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com>
Date: Sat, 9 Nov 2024 02:01:18 -0400
Subject: [PATCH] feat: rip out toolcache so we can publish to pypi (#116)

---
 .github/workflows/deploy_docs.yaml |  2 +-
 README.md                          |  7 +---
 eth_portfolio/_cache.py            | 57 +++++++++++++++++++++++-------
 3 files changed, 47 insertions(+), 19 deletions(-)

diff --git a/.github/workflows/deploy_docs.yaml b/.github/workflows/deploy_docs.yaml
index f0fe261e..eff155d1 100644
--- a/.github/workflows/deploy_docs.yaml
+++ b/.github/workflows/deploy_docs.yaml
@@ -28,7 +28,7 @@ jobs:
           pip install wheel
           pip install --no-build-isolation "Cython<3" "pyyaml==5.4.1"
           pip install -r requirements.txt
-          pip install sphinx sphinx-rtd-theme git+https://github.com/BobTheBuidler/toolcache@patch-1
+          pip install sphinx sphinx-rtd-theme
 
       - name: Setup brownie networks
         run: |
diff --git a/README.md b/README.md
index 0fc5b5ab..4f19eae8 100644
--- a/README.md
+++ b/README.md
@@ -7,12 +7,7 @@ Use `eth-portfolio` to output information about your portfolio in a streamlined,
 - This lib is still a WIP and the provided API is subject to change without notice.
 
 ### INSTALLATION:
-- You should start with a fresh virtual environment
-- First, install my modified version of toolcache since I am not permitted to include a git hash in a requirements.txt file
-```
-pip install git+https://github.com/BobTheBuidler/toolcache@patch-1
-```
-- Next, install the lib
+- You should start with a fresh virtual environment, and then just...
 ```
 pip install git+https://github.com/BobTheBuidler/eth-portfolio
 ```
diff --git a/eth_portfolio/_cache.py b/eth_portfolio/_cache.py
index 051b5dbb..57808224 100644
--- a/eth_portfolio/_cache.py
+++ b/eth_portfolio/_cache.py
@@ -1,35 +1,68 @@
+import contextlib
 import functools
+import hashlib
+import inspect
 import os
-from inspect import iscoroutinefunction
+import pickle
+from typing import Any
 
-import toolcache
+import a_sync
+import aiofiles
+import aiofiles.os
 from a_sync._typing import AnyFn, P, T
 from brownie import chain
 
-cache_base_path = f"./cache/{chain.id}/"
+BASE_PATH = f"./cache/{chain.id}/"
 
 
 def cache_to_disk(fn: AnyFn[P, T]) -> AnyFn[P, T]:
     module = fn.__module__.replace(".", "/")
-    cache_path_for_fn = cache_base_path + module + "/" + fn.__name__
+    cache_path_for_fn = BASE_PATH + module + "/" + fn.__name__
 
-    # Ensure the cache dir exists
-    os.makedirs(cache_path_for_fn, exist_ok=True)
+    def get_cache_file_path(args, kwargs):
+        # Create a unique filename based on the function arguments
+        key = hashlib.md5(pickle.dumps((args, sorted(kwargs.items())))).hexdigest()
+        return os.path.join(cache_path_for_fn, f"{key}.json")
 
-    cache_decorator = toolcache.cache("disk", cache_dir=cache_path_for_fn)
+    os.makedirs(cache_path_for_fn, exist_ok=True)
 
-    if iscoroutinefunction(fn):
+    if inspect.iscoroutinefunction(fn):
 
-        @cache_decorator
         @functools.wraps(fn)
         async def disk_cache_wrap(*args: P.args, **kwargs: P.kwargs) -> T:
-            return await fn(*args, **kwargs)
+            cache_path = get_cache_file_path(args, kwargs)
+            if await aiofiles.os.path.exists(cache_path):
+                async with aiofiles.open(cache_path, "rb") as f:
+                    with contextlib.suppress(EOFError):
+                        return pickle.loads(await f.read())
+            async_result: T = await fn(*args, **kwargs)
+            _cache_write(cache_path, async_result)
+            return async_result
 
     else:
 
-        @cache_decorator
         @functools.wraps(fn)
         def disk_cache_wrap(*args: P.args, **kwargs: P.kwargs) -> T:
-            return fn(*args, **kwargs)  # type: ignore [return-value]
+            cache_path = get_cache_file_path(args, kwargs)
+            if os.path.exists(cache_path):
+                with open(cache_path, "rb") as f:
+                    with contextlib.suppress(EOFError):
+                        return pickle.load(f)
+
+            sync_result: T = fn(*args, **kwargs)  # type: ignore [assignment, return-value]
+            _cache_write(cache_path, sync_result)
+            return sync_result
 
     return disk_cache_wrap
+
+
+def _cache_write(cache_path: str, result: Any) -> None:
+    a_sync.create_task(
+        coro=__cache_write(cache_path, result),
+        skip_gc_until_done=True,
+    )
+
+
+async def __cache_write(cache_path: str, result: Any) -> None:
+    async with aiofiles.open(cache_path, "wb") as f:
+        await f.write(pickle.dumps(result))