from dataclasses import dataclass, field, replace
from typing import Any
from inspect import signature
from collections.abc import Mapping
from . import digest
@dataclass
[docs]
class Call:
"""
Represents a function call, capturing its name, arguments, and keyword arguments.
`module` and `version` can be optionally set to be included in the hash of the call.
`version` should be a plain integer and monotonically increase. Each different version will completely change the
hash of the call, invalidating previously cached results.
"""
[docs]
arguments: dict[str, Any]
[docs]
module: str | None = None
[docs]
version: int | None = None
[docs]
code_digest: str | None = None
@classmethod
[docs]
def from_call(cls, func, *args, **kwargs):
# Normalize arguments using function signature
sig = signature(func)
bound = sig.bind(*args, **kwargs)
bound.apply_defaults()
arguments = dict(bound.arguments)
# Preserve declared parameter order via bound.arguments (OrderedDict)
call = cls(func.__name__, arguments)
if hasattr(func, "__version__"):
call.version = func.__version__
if hasattr(func, "__module__"):
call.module = func.__module__
if hasattr(func, "__code__"):
call.code_digest = digest.digest(func.__code__)
return call
[docs]
def to_lookup_key(self):
# Iterate explicitly in the preserved parameter order; do not sort
arg_pairs = tuple(self.arguments.items())
call = replace(self, arguments=arg_pairs, metadata=None, result=None)
return digest.digest(call)
[docs]
def matches(self, other: 'Call | LazyCall') -> bool:
"""Check if this call matches another call, treating None as a wildcard in this object."""
def none_or_equal(a, b):
if a is None:
return True
# Use digest to handle both raw values and Digest objects consistently
return digest.digest(a) == digest.digest(b)
if not none_or_equal(self.name, other.name):
return False
if not none_or_equal(self.module, other.module):
return False
if not none_or_equal(self.version, other.version):
return False
if not none_or_equal(self.result, other.result):
return False
if self.arguments is not None:
for k, v in self.arguments.items():
if k not in other.arguments:
return False
if not none_or_equal(v, other.arguments[k]):
return False
if self.metadata:
for mname, filters in self.metadata.items():
data = other.metadata.get(mname)
if data is None:
return False
for kk, vv in (filters or {}).items():
if vv is None:
if kk not in data:
return False
else:
if data.get(kk) != vv:
return False
return True
class LazyArguments(Mapping):
def __init__(self, cache, arg_digests):
self._cache = cache
self._arg_digests = arg_digests
def __getitem__(self, key):
return self._cache._handle_args_load(self._arg_digests[key])
def __iter__(self):
return iter(self._arg_digests)
def __len__(self):
return len(self._arg_digests)
def __repr__(self):
return f"LazyArguments({self._arg_digests!r})"
def __digest__(self):
# Ensuring that LazyArguments digests identically to a dict of the same values.
# Since self._arg_digests are already Digests, and digest(Digest(X)) == X,
# this will match a dict of raw values because digest(val) == X.
return digest.digest(self._arg_digests)
@dataclass(frozen=True)
[docs]
class LazyCall:
[docs]
_arguments: dict[str, Any]
[docs]
_cache: Any = field(repr=False, compare=False)
[docs]
module: str | None = None
[docs]
version: int | None = None
[docs]
code_digest: str | None = None
@property
[docs]
def arguments(self):
return LazyArguments(self._cache, self._arguments)
@property
[docs]
def result(self):
return self._cache.load_value(self._result)
[docs]
def to_lookup_key(self) -> str:
# Reconstruct a Call object to ensure identical key calculation
c = Call(
name=self.name,
arguments=self._arguments,
metadata=self.metadata,
module=self.module,
version=self.version,
code_digest=self.code_digest,
result=None
)
return c.to_lookup_key()
[docs]
def fetch(self) -> Call:
"""Reconstruct a full Call object by loading all values from the cache."""
return Call(
name=self.name,
arguments=dict(self.arguments),
metadata=self.metadata,
module=self.module,
version=self.version,
code_digest=self.code_digest,
result=self.result
)
[docs]
def __digest__(self):
# Reconstruct a Call object to ensure identical digest calculation
c = Call(
name=self.name,
arguments=self._arguments,
metadata=self.metadata,
module=self.module,
version=self.version,
code_digest=self.code_digest,
result=self._result
)
return digest.digest(c)
AnyQueryType = None | digest.Digest | Any
StrQueryType = None | digest.Digest | str
@dataclass
[docs]
class QueryCall:
[docs]
name: StrQueryType = None
[docs]
arguments: dict[str, AnyQueryType] | None = None
[docs]
module: str | None = None
[docs]
version: int | None = None
[docs]
code_digest: digest.Digest | None = None
[docs]
result: AnyQueryType = None
@classmethod
[docs]
def from_call(cls, func, *args, **kwargs):
# Normalize arguments using function signature
sig = signature(func)
bound = sig.bind_partial(*args, **kwargs)
bound.apply_defaults()
# missing arguments are set to None
arguments = {name: bound.arguments.get(name) for name in sig.parameters}
call = cls(func.__name__, arguments)
if hasattr(func, "__version__"):
call.version = func.__version__
if hasattr(func, "__module__"):
call.module = func.__module__
return call
[docs]
def matches(self, other: 'Call | LazyCall') -> bool:
"""Check if this call matches another call, treating None as a wildcard in this object."""
def none_or_equal(a, b):
if a is None:
return True
# Use digest to handle both raw values and Digest objects consistently
return digest.digest(a) == digest.digest(b)
if not none_or_equal(self.name, other.name):
return False
if not none_or_equal(self.module, other.module):
return False
if not none_or_equal(self.version, other.version):
return False
if not none_or_equal(self.code_digest, other.code_digest):
return False
if not none_or_equal(self.result, other.result):
return False
if self.arguments is not None:
for k, v in self.arguments.items():
if k not in other.arguments:
return False
if not none_or_equal(v, other.arguments[k]):
return False
if self.metadata:
for mname, filters in self.metadata.items():
data = other.metadata.get(mname)
if data is None:
return False
for kk, vv in (filters or {}).items():
if vv is None:
if kk not in data:
return False
else:
if data.get(kk) != vv:
return False
return True
[docs]
AnyCall = Call | LazyCall
__all__ = [
"Call",
"LazyCall",
"QueryCall",
"AnyCall"
]