speedy-utils 1.1.9__py3-none-any.whl → 1.1.11__py3-none-any.whl

speedy_utils/common/utils_cache.py

@@ -6,8 +6,15 @@ import os
 import os.path as osp
 import pickle
 import uuid
+import weakref
 from threading import Lock
-from typing import Any, Awaitable, Callable, Literal, TypeVar
+from typing import Any, Awaitable, Callable, Literal, Optional, TypeVar, overload
+
+try:
+    # Python 3.10+
+    from typing import ParamSpec
+except ImportError:  # pragma: no cover
+    from typing_extensions import ParamSpec  # type: ignore
 
 import cachetools
 import pandas as pd
@@ -18,155 +25,366 @@ from pydantic import BaseModel
 from speedy_utils.common.utils_io import dump_json_or_pickle, load_json_or_pickle
 from speedy_utils.common.utils_misc import mkdir_or_exist
 
-SPEED_CACHE_DIR = osp.join(osp.expanduser("~"), ".cache/speedy_cache")
-LRU_MEM_CACHE = cachetools.LRUCache(maxsize=128_000)
+# --------------------------------------------------------------------------------------
+# Defaults / Globals
+# --------------------------------------------------------------------------------------
 
-thread_locker = Lock()
+SPEED_CACHE_DIR = osp.join(osp.expanduser("~"), ".cache/speedy_cache")
 
-# Add two locks for thread-safe cache access
+# Thread locks for safety
 disk_lock = Lock()
 mem_lock = Lock()
 
-# Add async-specific types
-T = TypeVar('T')
-AsyncFunc = Callable[..., Awaitable[T]]
+# Quick identifier cache for big objects that support weakref
+# (prevents recomputing expensive keys for the same object instance)
+_QUICK_ID_MAP: "weakref.WeakKeyDictionary[Any, str]" = weakref.WeakKeyDictionary()
 
+# Per-function memory caches (so different functions can have different LRU sizes)
+_MEM_CACHES: "weakref.WeakKeyDictionary[Callable[..., Any], cachetools.LRUCache]" = (
+    weakref.WeakKeyDictionary()
+)
 
-def compute_func_id(func, args, kwargs, ignore_self, keys):
-    func_source = get_source(func)
-    if keys:
-        arg_spec = inspect.getfullargspec(func).args
-        used_args = {arg_spec[i]: arg for i, arg in enumerate(args)}
-        used_args.update(kwargs)
-        values = [used_args[k] for k in keys if k in used_args]
-        if not values:
-            raise ValueError(f"Keys {keys} not found in function arguments")
-        param_hash = identify(values)
-        dir_path = f"{func.__name__}_{identify(func_source)}"
-        key_id = f"{'_'.join(keys)}_{param_hash}.pkl"
-        return func_source, dir_path, key_id
+# Backward-compat global symbol (internal only; not exported)
+LRU_MEM_CACHE = cachetools.LRUCache(maxsize=256)
 
-    if (
-        inspect.getfullargspec(func).args
-        and inspect.getfullargspec(func).args[0] == "self"
-        and ignore_self
-    ):
-        fid = (func_source, args[1:], kwargs)
-    else:
-        fid = (func_source, args, kwargs)
-    return func_source, "funcs", f"{identify(fid)}.pkl"
+# Typing helpers
+P = ParamSpec("P")
+R = TypeVar("R")
+AsyncFunc = Callable[P, Awaitable[R]]
+
+# --------------------------------------------------------------------------------------
+# Utilities
+# --------------------------------------------------------------------------------------
 
 
 def fast_serialize(x: Any) -> bytes:
+    """Serialize x quickly; JSON if possible (stable), else pickle."""
     try:
-        return json.dumps(x, sort_keys=True).encode("utf-8")
+        return json.dumps(x, sort_keys=True, default=str).encode("utf-8")
     except (TypeError, ValueError):
         return pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
 
 
-def identify(obj: Any, depth=0, max_depth=2) -> str:
+def identify_uuid(x: Any) -> str:
+    data = fast_serialize(x)
+    hash_obj = xxhash.xxh128(data, seed=0)
+    return str(uuid.UUID(bytes=hash_obj.digest()))
+
+
+def get_source(func: Callable[..., Any]) -> str:
+    """Minified function source; falls back to module + qualname for builtins/lambdas."""
+    try:
+        code = inspect.getsource(func)
+    except OSError:
+        # source not available (e.g., builtins, some C extensions)
+        mod = getattr(func, "__module__", "unknown")
+        qn = getattr(func, "__qualname__", getattr(func, "__name__", "unknown"))
+        code = f"{mod}.{qn}"
+    # normalize whitespace to make it stable
+    for r in (" ", "\n", "\t", "\r"):
+        code = code.replace(r, "")
+    return code
+
+
+def _try_get_quick_id(obj: Any) -> Optional[str]:
+    """Return a quick identifier if obj is weakref-able and cached."""
+    try:
+        return _QUICK_ID_MAP.get(obj)  # type: ignore[arg-type]
+    except TypeError:
+        # not weakref-able (e.g., list/dict); cannot use WeakKeyDictionary
+        return None
+
+
+def _try_store_quick_id(obj: Any, ident: str) -> None:
+    """Store quick identifier if obj is weakref-able."""
+    try:
+        _QUICK_ID_MAP[obj] = ident  # type: ignore[index]
+    except TypeError:
+        # not weakref-able
+        pass
+
+
+def identify(obj: Any, depth: int = 0, max_depth: int = 2) -> str:
+    """
+    Produce a stable, content-based identifier string for arbitrary Python objects.
+    Includes a quick path using a weakref cache for large, user-defined objects.
+    """
+    # Quick-path for user-defined objects (weakref-able)
+    if depth == 0:
+        quick = _try_get_quick_id(obj)
+        if quick is not None:
+            return quick
+
     if isinstance(obj, (list, tuple)):
         x = [identify(x, depth + 1, max_depth) for x in obj]
         x = "\n".join(x)
-        return identify(x, depth + 1, max_depth)
-    # is pandas row or dict
+        out = identify(x, depth + 1, max_depth)
+        if depth == 0:
+            _try_store_quick_id(obj, out)
+        return out
     elif isinstance(obj, (pd.DataFrame, pd.Series)):
         x = str(obj.to_dict())
-        return identify(x, depth + 1, max_depth)
+        out = identify(x, depth + 1, max_depth)
+        if depth == 0:
+            _try_store_quick_id(obj, out)
+        return out
     elif hasattr(obj, "__code__"):
-        return identify(get_source(obj), depth + 1, max_depth)
+        out = identify(get_source(obj), depth + 1, max_depth)
+        if depth == 0:
+            _try_store_quick_id(obj, out)
+        return out
     elif isinstance(obj, BaseModel):
-        obj = obj.model_dump()
-        return identify(obj, depth + 1, max_depth)
+        out = identify(obj.model_dump(), depth + 1, max_depth)
+        if depth == 0:
+            _try_store_quick_id(obj, out)
+        return out
     elif isinstance(obj, dict):
         ks = sorted(obj.keys())
         vs = [identify(obj[k], depth + 1, max_depth) for k in ks]
-        return identify([ks, vs], depth + 1, max_depth)
+        out = identify([ks, vs], depth + 1, max_depth)
+        if depth == 0:
+            _try_store_quick_id(obj, out)
+        return out
     elif obj is None:
-        return identify("None", depth + 1, max_depth)
+        out = identify("None", depth + 1, max_depth)
+        if depth == 0:
+            _try_store_quick_id(obj, out)
+        return out
     else:
-        # primitive_types = [int, float, str, bool]
-        # if not type(obj) in primitive_types:
-        #     logger.warning(f"Unknown type: {type(obj)}")
-        return xxhash.xxh64_hexdigest(fast_serialize(obj), seed=0)
+        # primitives / everything else
+        out = xxhash.xxh64_hexdigest(fast_serialize(obj), seed=0)
+        if depth == 0:
+            _try_store_quick_id(obj, out)
+        return out
+
+
+def _build_named_keys(
+    func: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    keys: list[str],
+) -> list[Any]:
+    """Extract named parameters in order from args/kwargs for keying."""
+    arg_spec = inspect.getfullargspec(func).args
+    used_args = {arg_spec[i]: arg for i, arg in enumerate(args[: len(arg_spec)])}
+    used_args.update(kwargs)
+    values = [used_args[k] for k in keys if k in used_args]
+    if not values:
+        raise ValueError(f"Keys {keys} not found in function arguments")
+    return values
+
+
+def _compute_cache_components(
+    func: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    ignore_self: bool,
+    keys: Optional[list[str]],
+    key_fn: Optional[Callable[..., Any]],
+):
+    """
+    Return (func_source, sub_dir, key_id) for disk paths and memory keying.
+    - If key_fn provided, it determines the cache key content.
+    - Else if keys list provided, use those argument names.
+    - Else use full (args, kwargs), optionally ignoring 'self' for methods.
+    """
+    func_source = get_source(func)
 
+    # Custom key function (most explicit & fastest when user knows what's important)
+    if key_fn is not None:
+        try:
+            custom_val = key_fn(*args, **kwargs)
+        except Exception as e:
+            raise ValueError(f"key function for {func.__name__} raised: {e}") from e
+        sub_dir = "custom"
+        key_id = f"{identify(custom_val)}.pkl"
+        return func_source, sub_dir, key_id
 
-def identify_uuid(x: Any) -> str:
-    data = fast_serialize(x)
-    hash_obj = xxhash.xxh128(data, seed=0)
-    return str(uuid.UUID(bytes=hash_obj.digest()))
+    # Named keys (back-compat)
+    if keys:
+        values = _build_named_keys(func, args, kwargs, keys)
+        param_hash = identify(values)
+        dir_path = f"{func.__name__}_{identify(func_source)}"
+        key_id = f"{'_'.join(keys)}_{param_hash}.pkl"
+        return func_source, dir_path, key_id
 
+    # Default: full argument identity (optionally ignoring 'self')
+    if (
+        inspect.getfullargspec(func).args
+        and inspect.getfullargspec(func).args[0] == "self"
+        and ignore_self
+    ):
+        fid = (func_source, args[1:], kwargs)
+    else:
+        fid = (func_source, args, kwargs)
 
-def get_source(func):
-    code = inspect.getsource(func)
-    for r in [" ", "\n", "\t", "\r"]:
-        code = code.replace(r, "")
-    return code
+    return func_source, "funcs", f"{identify(fid)}.pkl"
 
 
-def _disk_memoize(func, keys, cache_dir, ignore_self, verbose):
+def _mem_cache_for(func: Callable[..., Any], size: int) -> cachetools.LRUCache:
+    """Get or create a per-function LRU cache with the given size."""
+    # Keep a per-function cache to avoid cross-talk of maxsize across functions
+    with mem_lock:
+        cache = _MEM_CACHES.get(func)
+        if cache is None or cache.maxsize != size:
+            cache = cachetools.LRUCache(maxsize=size)
+            _MEM_CACHES[func] = cache
+    # Keep global symbol backwards-compatible internally
+    global LRU_MEM_CACHE
+    LRU_MEM_CACHE = cache
+    return cache
+
+
+# --------------------------------------------------------------------------------------
+# Memory-only memoize (sync / async)
+# --------------------------------------------------------------------------------------
+
+
+def _memory_memoize(
+    func: Callable[P, R],
+    size: int,
+    keys: Optional[list[str]],
+    ignore_self: bool,
+    key_fn: Optional[Callable[..., Any]],
+) -> Callable[P, R]:
+    mem_cache = _mem_cache_for(func, size)
+
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        func_source, sub_dir, key_id = _compute_cache_components(
+            func, args, kwargs, ignore_self, keys, key_fn
+        )
+        name = identify((func_source, sub_dir, key_id))
+
+        with mem_lock:
+            if name in mem_cache:
+                return mem_cache[name]  # type: ignore[return-value]
+
+        result = func(*args, **kwargs)
+
+        with mem_lock:
+            if name not in mem_cache:
+                mem_cache[name] = result  # type: ignore[index]
+        return result
+
+    return wrapper
+
+
+def _async_memory_memoize(
+    func: AsyncFunc[P, R],
+    size: int,
+    keys: Optional[list[str]],
+    ignore_self: bool,
+    key_fn: Optional[Callable[..., Any]],
+) -> AsyncFunc[P, R]:
+    mem_cache = _mem_cache_for(func, size)
+
+    # Avoid duplicate in-flight computations for the same key
+    inflight: dict[str, asyncio.Task[R]] = {}
+    alock = asyncio.Lock()
+
+    @functools.wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        func_source, sub_dir, key_id = _compute_cache_components(
+            func, args, kwargs, ignore_self, keys, key_fn
+        )
+        name = identify((func_source, sub_dir, key_id))
+
+        async with alock:
+            if name in mem_cache:
+                return mem_cache[name]  # type: ignore[return-value]
+            task = inflight.get(name)
+            if task is None:
+                task = asyncio.create_task(func(*args, **kwargs))  # type: ignore[arg-type]
+                inflight[name] = task
+
         try:
-            # Compute cache path as before
-            func_source, sub_dir, key_id = compute_func_id(
-                func, args, kwargs, ignore_self, keys
+            result = await task
+        finally:
+            async with alock:
+                inflight.pop(name, None)
+
+        with mem_lock:
+            mem_cache[name] = result  # type: ignore[index]
+        return result
+
+    return wrapper
+
+
+# --------------------------------------------------------------------------------------
+# Disk-only memoize (sync / async)
+# --------------------------------------------------------------------------------------
+
+
+def _disk_memoize(
+    func: Callable[P, R],
+    keys: Optional[list[str]],
+    cache_dir: str,
+    ignore_self: bool,
+    verbose: bool,
+    key_fn: Optional[Callable[..., Any]],
+) -> Callable[P, R]:
+    @functools.wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        try:
+            func_source, sub_dir, key_id = _compute_cache_components(
+                func, args, kwargs, ignore_self, keys, key_fn
             )
-            if func_source is None:
-                return func(*args, **kwargs)
             if sub_dir == "funcs":
                 cache_path = osp.join(cache_dir, sub_dir, func.__name__, key_id)
             else:
                 cache_path = osp.join(cache_dir, sub_dir, key_id)
             mkdir_or_exist(osp.dirname(cache_path))
 
-            # First check with disk lock
             with disk_lock:
                 if osp.exists(cache_path):
-                    # logger.debug(f"Cache HIT for {func.__name__}, key={cache_path}")
                     try:
                         return load_json_or_pickle(cache_path)
                     except Exception as e:
                         if osp.exists(cache_path):
                             os.remove(cache_path)
-                        logger.opt(depth=1).warning(
-                            f"Error loading cache: {str(e)[:100]}, continue to recompute"
-                        )
+                        if verbose:
+                            logger.opt(depth=1).warning(
+                                f"Error loading cache: {str(e)[:100]}, recomputing"
+                            )
 
             result = func(*args, **kwargs)
 
-            # Write result under disk lock to avoid race conditions
             with disk_lock:
                 if not osp.exists(cache_path):
                     dump_json_or_pickle(result, cache_path)
             return result
         except Exception as e:
-            logger.opt(depth=1).warning(
-                f"Failed to cache {func.__name__}: {e}, continue to recompute without cache"
-            )
+            if verbose:
+                logger.opt(depth=1).warning(
+                    f"Failed to cache {func.__name__}: {e}, executing without cache"
+                )
             return func(*args, **kwargs)
 
     return wrapper
 
 
-def _async_disk_memoize(func, keys, cache_dir, ignore_self, verbose):
+def _async_disk_memoize(
+    func: AsyncFunc[P, R],
+    keys: Optional[list[str]],
+    cache_dir: str,
+    ignore_self: bool,
+    verbose: bool,
+    key_fn: Optional[Callable[..., Any]],
+) -> AsyncFunc[P, R]:
     @functools.wraps(func)
-    async def wrapper(*args, **kwargs):
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
         try:
-            # Compute cache path as before
-            func_source, sub_dir, key_id = compute_func_id(
-                func, args, kwargs, ignore_self, keys
+            func_source, sub_dir, key_id = _compute_cache_components(
+                func, args, kwargs, ignore_self, keys, key_fn
             )
-            if func_source is None:
-                return await func(*args, **kwargs)
             if sub_dir == "funcs":
                 cache_path = osp.join(cache_dir, sub_dir, func.__name__, key_id)
             else:
                 cache_path = osp.join(cache_dir, sub_dir, key_id)
             mkdir_or_exist(osp.dirname(cache_path))
 
-            # First check with disk lock (run in thread to avoid blocking)
-            def check_cache():
+            def check_cache() -> Optional[R]:
                 with disk_lock:
                     if osp.exists(cache_path):
                         try:
@@ -174,12 +392,12 @@ def _async_disk_memoize(func, keys, cache_dir, ignore_self, verbose):
                         except Exception as e:
                             if osp.exists(cache_path):
                                 os.remove(cache_path)
-                            logger.opt(depth=1).warning(
-                                f"Error loading cache: {str(e)[:100]}, continue to recompute"
-                            )
+                            if verbose:
+                                logger.opt(depth=1).warning(
+                                    f"Error loading cache: {str(e)[:100]}, recomputing"
+                                )
                     return None
-            
-            # Run cache check in thread pool to avoid blocking
+
             loop = asyncio.get_event_loop()
             cached_result = await loop.run_in_executor(None, check_cache)
             if cached_result is not None:
@@ -187,157 +405,117 @@ def _async_disk_memoize(func, keys, cache_dir, ignore_self, verbose):
 
             result = await func(*args, **kwargs)
 
-            # Write result under disk lock (run in thread to avoid blocking)
-            def write_cache():
+            def write_cache() -> None:
                 with disk_lock:
                     if not osp.exists(cache_path):
                         dump_json_or_pickle(result, cache_path)
-            
+
             await loop.run_in_executor(None, write_cache)
             return result
         except Exception as e:
-            logger.opt(depth=1).warning(
-                f"Failed to cache {func.__name__}: {e}, continue to recompute without cache"
-            )
+            if verbose:
+                logger.opt(depth=1).warning(
+                    f"Failed to cache {func.__name__}: {e}, executing without cache"
+                )
             return await func(*args, **kwargs)
 
     return wrapper
 
 
-def _memory_memoize(func, size, keys, ignore_self):
-    global LRU_MEM_CACHE
-    if LRU_MEM_CACHE.maxsize != size:
-        LRU_MEM_CACHE = cachetools.LRUCache(maxsize=size)
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        func_source, sub_dir, key_id = compute_func_id(
-            func, args, kwargs, ignore_self, keys
-        )
-        if func_source is None:
-            return func(*args, **kwargs)
-        name = identify((func_source, sub_dir, key_id))
-
-        if not hasattr(func, "_mem_cache"):
-            func._mem_cache = LRU_MEM_CACHE
-
-        with mem_lock:
-            if name in func._mem_cache:
-                # logger.debug(f"Cache HIT (memory) for {func.__name__}, key={name}")
-                return func._mem_cache[name]
+# --------------------------------------------------------------------------------------
+# Memory+Disk (sync / async)
+# --------------------------------------------------------------------------------------
 
-        result = func(*args, **kwargs)
 
-        with mem_lock:
-            if name not in func._mem_cache:
-                func._mem_cache[name] = result
-        return result
-
-    return wrapper
-
-
-def _async_memory_memoize(func, size, keys, ignore_self):
-    global LRU_MEM_CACHE
-    if LRU_MEM_CACHE.maxsize != size:
-        LRU_MEM_CACHE = cachetools.LRUCache(maxsize=size)
+def both_memoize(
+    func: Callable[P, R],
+    keys: Optional[list[str]],
+    cache_dir: str,
+    ignore_self: bool,
+    size: int,
+    key_fn: Optional[Callable[..., Any]],
+) -> Callable[P, R]:
+    mem_cache = _mem_cache_for(func, size)
 
     @functools.wraps(func)
-    async def wrapper(*args, **kwargs):
-        func_source, sub_dir, key_id = compute_func_id(
-            func, args, kwargs, ignore_self, keys
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        func_source, sub_dir, key_id = _compute_cache_components(
+            func, args, kwargs, ignore_self, keys, key_fn
         )
-        if func_source is None:
-            return await func(*args, **kwargs)
-        name = identify((func_source, sub_dir, key_id))
-
-        if not hasattr(func, "_mem_cache"):
-            func._mem_cache = LRU_MEM_CACHE
-
-        with mem_lock:
-            if name in func._mem_cache:
-                # logger.debug(f"Cache HIT (memory) for {func.__name__}, key={name}")
-                return func._mem_cache[name]
-
-        result = await func(*args, **kwargs)
-
-        with mem_lock:
-            if name not in func._mem_cache:
-                func._mem_cache[name] = result
-        return result
-
-    return wrapper
-
-
-def both_memoize(func, keys, cache_dir, ignore_self):
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        func_source, sub_dir, key_id = compute_func_id(
-            func, args, kwargs, ignore_self, keys
-        )
-        if func_source is None:
-            return func(*args, **kwargs)
-
         mem_key = identify((func_source, sub_dir, key_id))
-        if not hasattr(func, "_mem_cache"):
-            func._mem_cache = LRU_MEM_CACHE
 
+        # Memory first
         with mem_lock:
-            if mem_key in func._mem_cache:
-                # logger.debug(f"Cache HIT (memory) for {func.__name__}, key={mem_key}")
-                return func._mem_cache[mem_key]
+            if mem_key in mem_cache:
+                return mem_cache[mem_key]  # type: ignore[return-value]
 
+        # Disk next
         if sub_dir == "funcs":
             cache_path = osp.join(cache_dir, sub_dir, func.__name__, key_id)
         else:
             cache_path = osp.join(cache_dir, sub_dir, key_id)
         mkdir_or_exist(osp.dirname(cache_path))
 
+        disk_result: Optional[R] = None
         with disk_lock:
             if osp.exists(cache_path):
-                # logger.debug(f"Cache HIT (disk) for {func.__name__}, key={cache_path}")
-                result = load_json_or_pickle(cache_path)
-                with mem_lock:
-                    func._mem_cache[mem_key] = result
-                return result
-        # logger.debug(f"Cache MISS for {func.__name__}, key={cache_path}")
-        result = func(*args, **kwargs)
+                try:
+                    disk_result = load_json_or_pickle(cache_path)
+                except Exception:
+                    if osp.exists(cache_path):
+                        os.remove(cache_path)
+                    disk_result = None
 
+        if disk_result is not None:
+            with mem_lock:
+                mem_cache[mem_key] = disk_result  # type: ignore[index]
+            return disk_result
+
+        # Miss: compute, then write both
+        result = func(*args, **kwargs)
         with disk_lock:
             if not osp.exists(cache_path):
                 dump_json_or_pickle(result, cache_path)
         with mem_lock:
-            func._mem_cache[mem_key] = result
+            mem_cache[mem_key] = result  # type: ignore[index]
         return result
 
     return wrapper
 
 
-def _async_both_memoize(func, keys, cache_dir, ignore_self):
+def _async_both_memoize(
+    func: AsyncFunc[P, R],
+    keys: Optional[list[str]],
+    cache_dir: str,
+    ignore_self: bool,
+    size: int,
+    key_fn: Optional[Callable[..., Any]],
+) -> AsyncFunc[P, R]:
+    mem_cache = _mem_cache_for(func, size)
+
+    inflight: dict[str, asyncio.Task[R]] = {}
+    alock = asyncio.Lock()
+
     @functools.wraps(func)
-    async def wrapper(*args, **kwargs):
-        func_source, sub_dir, key_id = compute_func_id(
-            func, args, kwargs, ignore_self, keys
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        func_source, sub_dir, key_id = _compute_cache_components(
+            func, args, kwargs, ignore_self, keys, key_fn
         )
-        if func_source is None:
-            return await func(*args, **kwargs)
-
         mem_key = identify((func_source, sub_dir, key_id))
-        if not hasattr(func, "_mem_cache"):
-            func._mem_cache = LRU_MEM_CACHE
 
-        with mem_lock:
-            if mem_key in func._mem_cache:
-                # logger.debug(f"Cache HIT (memory) for {func.__name__}, key={mem_key}")
-                return func._mem_cache[mem_key]
+        # Memory
+        async with alock:
+            if mem_key in mem_cache:
+                return mem_cache[mem_key]  # type: ignore[return-value]
 
+        # Disk
         if sub_dir == "funcs":
             cache_path = osp.join(cache_dir, sub_dir, func.__name__, key_id)
         else:
             cache_path = osp.join(cache_dir, sub_dir, key_id)
         mkdir_or_exist(osp.dirname(cache_path))
 
-        # Check disk cache in thread pool to avoid blocking
-        def check_disk_cache():
+        def check_disk_cache() -> Optional[R]:
             with disk_lock:
                 if osp.exists(cache_path):
                     return load_json_or_pickle(cache_path)
@@ -345,150 +523,142 @@ def _async_both_memoize(func, keys, cache_dir, ignore_self):
 
         loop = asyncio.get_event_loop()
         disk_result = await loop.run_in_executor(None, check_disk_cache)
-        
+
         if disk_result is not None:
             with mem_lock:
-                func._mem_cache[mem_key] = disk_result
+                mem_cache[mem_key] = disk_result  # type: ignore[index]
             return disk_result
 
-        # logger.debug(f"Cache MISS for {func.__name__}, key={cache_path}")
-        result = await func(*args, **kwargs)
+        # Avoid duplicate async work for same key
+        async with alock:
+            task = inflight.get(mem_key)
+            if task is None:
+                task = asyncio.create_task(func(*args, **kwargs))  # type: ignore[arg-type]
+                inflight[mem_key] = task
+
+        try:
+            result = await task
+        finally:
+            async with alock:
+                inflight.pop(mem_key, None)
 
-        # Write to disk in thread pool to avoid blocking
-        def write_disk_cache():
+        def write_disk_cache() -> None:
             with disk_lock:
                 if not osp.exists(cache_path):
                     dump_json_or_pickle(result, cache_path)
 
         await loop.run_in_executor(None, write_disk_cache)
-        
+
         with mem_lock:
-            func._mem_cache[mem_key] = result
+            mem_cache[mem_key] = result  # type: ignore[index]
         return result
 
     return wrapper
 
 
-def memoize(
-    _func=None,
-    *,
-    keys=None,
-    cache_dir=SPEED_CACHE_DIR,
-    cache_type: Literal["memory", "disk", "both"] = "disk",
-    size=10240,
-    ignore_self=True,
-    verbose=False,
-):
-    if "~/" in cache_dir:
-        cache_dir = osp.expanduser(cache_dir)
+# --------------------------------------------------------------------------------------
+# Public decorator (only export memoize)
+# --------------------------------------------------------------------------------------
 
-    def decorator(func):
-        # Check if function is async
-        is_async = inspect.iscoroutinefunction(func)
-        
-        if cache_type == "memory":
-            if is_async:
-                return _async_memory_memoize(
-                    func,
-                    size,
-                    keys,
-                    ignore_self,
-                )
-            return _memory_memoize(
-                func,
-                size,
-                keys,
-                ignore_self,
-            )
-        elif cache_type == "disk":
-            if is_async:
-                return _async_disk_memoize(
-                    func,
-                    keys,
-                    cache_dir,
-                    ignore_self,
-                    verbose,
-                )
-            return _disk_memoize(
-                func,
-                keys,
-                cache_dir,
-                ignore_self,
-                verbose,
-            )
-        
-        # cache_type == "both"
-        if is_async:
-            return _async_both_memoize(
-                func,
-                keys,
-                cache_dir,
-                ignore_self,
-            )
-        return both_memoize(
-            func,
-            keys,
-            cache_dir,
-            verbose,
-        )
 
-    # Handle both @memoize and @memoize() usage patterns
-    if _func is None:
-        return decorator
-    else:
-        return decorator(_func)
+@overload
+def memoize(
+    _func: Callable[P, R],
+    *,
+    keys: Optional[list[str]] = ...,
+    key: Optional[Callable[..., Any]] = ...,
+    cache_dir: str = ...,
+    cache_type: Literal["memory", "disk", "both"] = ...,
+    size: int = ...,
+    ignore_self: bool = ...,
+    verbose: bool = ...,
+) -> Callable[P, R]: ...
+@overload
+def memoize(
+    _func: Callable[P, Awaitable[R]],
+    *,
+    keys: Optional[list[str]] = ...,
+    key: Optional[Callable[..., Any]] = ...,
+    cache_dir: str = ...,
+    cache_type: Literal["memory", "disk", "both"] = ...,
+    size: int = ...,
+    ignore_self: bool = ...,
+    verbose: bool = ...,
+) -> Callable[P, Awaitable[R]]: ...
+@overload
+def memoize(
+    _func: None = ...,
+    *,
+    keys: Optional[list[str]] = ...,
+    key: Optional[Callable[..., Any]] = ...,
+    cache_dir: str = ...,
+    cache_type: Literal["memory", "disk", "both"] = ...,
+    size: int = ...,
+    ignore_self: bool = ...,
+    verbose: bool = ...,
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
 
 
-def amemoize(
-    _func=None,
+def memoize(
+    _func: Optional[Callable[P, Any]] = None,
     *,
-    keys: list[str] | None = None,
+    keys: Optional[list[str]] = None,
+    key: Optional[Callable[..., Any]] = None,
     cache_dir: str = SPEED_CACHE_DIR,
-    cache_type: Literal["memory", "disk", "both"] = "disk",
-    size: int = 10240,
+    cache_type: Literal["memory", "disk", "both"] = "both",
+    size: int = 256,
     ignore_self: bool = True,
     verbose: bool = False,
 ):
     """
-    Async-specific memoization decorator for coroutine functions.
-    
-    Args:
-        _func: The async function to memoize (when used without parentheses)
-        keys: Specific argument keys to use for cache key generation
-        cache_dir: Directory for disk cache storage
-        cache_type: Type of caching - "memory", "disk", or "both"  
-        size: Size of memory cache (for memory/both types)
-        ignore_self: Whether to ignore 'self' parameter in cache key
-        verbose: Enable verbose logging
-        
-    Returns:
-        Decorated async function with memoization
-        
-    Example:
-        @amemoize(cache_type="both")
-        async def my_async_func(x: int) -> str:
-            return str(x)
+    Universal memoizer that supports sync and async functions, preserves annotations
+    for Pylance via ParamSpec/TypeVar, and caches in memory + disk by default.
+
+    - keys: list of argument names to include in key (back-compat).
+    - key:  custom callable (*args, **kwargs) -> hashable/serializable object for keying.
+            Prefer this for performance on big inputs (e.g., key=lambda x: x.id).
+    - cache_dir: disk cache base directory (unlimited size).
+    - cache_type: "memory" | "disk" | "both" (default "both").
+    - size: memory LRU size per-function (default 256 items).
+    - ignore_self: ignore 'self' when building the default key for bound methods.
+    - verbose: enable warnings on cache load/write errors.
     """
     if "~/" in cache_dir:
         cache_dir = osp.expanduser(cache_dir)
+    from speedy_utils import timef
+
+    def decorator(func: Callable[P, Any]) -> Callable[P, Any]:
+        is_async = inspect.iscoroutinefunction(func)
+
+        # Apply timing decorator if verbose=True
+        target_func = timef(func) if verbose else func
 
-    def decorator(func):
-        # Ensure the function is actually async
-        if not inspect.iscoroutinefunction(func):
-            raise ValueError(f"amemoize can only be used with async functions. {func.__name__} is not async.")
-        
         if cache_type == "memory":
-            return _async_memory_memoize(func, size, keys, ignore_self)
-        elif cache_type == "disk":
-            return _async_disk_memoize(func, keys, cache_dir, ignore_self, verbose)
-        else:  # cache_type == "both"
-            return _async_both_memoize(func, keys, cache_dir, ignore_self)
+            if is_async:
+                return _async_memory_memoize(target_func, size, keys, ignore_self, key)  # type: ignore[return-value]
+            return _memory_memoize(target_func, size, keys, ignore_self, key)  # type: ignore[return-value]
+
+        if cache_type == "disk":
+            if is_async:
+                return _async_disk_memoize(
+                    target_func, keys, cache_dir, ignore_self, verbose, key
+                )  # type: ignore[return-value]
+            return _disk_memoize(
+                target_func, keys, cache_dir, ignore_self, verbose, key
+            )  # type: ignore[return-value]
+
+        # cache_type == "both"
+        if is_async:
+            return _async_both_memoize(
+                target_func, keys, cache_dir, ignore_self, size, key
+            )  # type: ignore[return-value]
+        return both_memoize(target_func, keys, cache_dir, ignore_self, size, key)  # type: ignore[return-value]
 
-    # Handle both @amemoize and @amemoize() usage patterns
+    # Support both @memoize and @memoize(...)
     if _func is None:
         return decorator
     else:
         return decorator(_func)
 
 
-__all__ = ["memoize", "identify", "identify_uuid", "amemoize"]
+__all__ = ["memoize", "identify"]