diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index 0b59b8223..7b32cb9c2 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -1,22 +1,42 @@ --- title: "*Memoization keys* & states" description: > - Customize how CocoIndex fingerprints memoized inputs via __coco_memo_key__ or - registered functions, and layer state validation (e.g. mtime, then content hash) - on top — plus NotMemoKeyable to opt out for stateful types. + Customize how CocoIndex fingerprints memoized inputs via __coco_memo_key__, + registered memo key functions, or stable type IDs, and layer state validation + (e.g. mtime, then content hash) on top. --- -As described in [Function — Change detection](../programming_guide/function#change-detection), CocoIndex detects [logic, input, and context changes](../programming_guide/function#change-detection) to decide whether a memo can be reused. Function arguments, [`deps`](../programming_guide/function#deps) values, and [context values](../programming_guide/context#change-detection) with `detect_change=True` are all fingerprinted through the same **data fingerprinting** pipeline. By default, most types are fingerprinted automatically. This page covers how to customize that pipeline — how objects are fingerprinted and validated: +As described in [Function — Change detection](../programming_guide/function#change-detection), CocoIndex detects [logic, input, and context changes](../programming_guide/function#change-detection) to decide whether a memo can be reused. By default, most types are fingerprinted automatically. + +Fingerprinting happens at these points: + +- **Function arguments** — when the function is called. +- **Change-detected context values** — when they are provided (for example, by `builder.provide()`), before `use_context()` reads them. +- **`deps` values** — once, when the `@coco.fn(..., deps=...)` or `@coco.fn.as_async(..., deps=...)` decorator is applied. + +This page covers how to customize that pipeline — how objects are fingerprinted and validated: - **Memoization keys** — how to control what CocoIndex uses as the fingerprint for your objects. - **Memo states** — how to add post-fingerprint validation to check freshness beyond simple equality. +- **Stable type IDs** — how to keep a type's memo namespace stable across refactors. + +:::tip[Choose the smallest mechanism] +Most values need no customization. CocoIndex already handles primitives, containers, dataclasses, Pydantic models, class objects, and other picklable objects. + +- Only one function should transform or ignore an argument → use [`memo_key=`](#override-at-the-call-site-with-memo_key). +- You control the type and every use should share the same key → define [`__coco_memo_key__`](#define-__coco_memo_key__-when-you-control-the-type). +- You don't control the type and every use should share the same key → [register a key function](#register-a-key-function-when-you-dont-control-the-type). +- A matching key still needs a freshness check, such as mtime or ETag → use a [memo state](#memo-state-validation). +- A class may move or be renamed while keeping the same memo namespace → use a [stable type ID](#use-stable-type-ids-across-refactors). +- A value must never participate in memoization → inherit from [`NotMemoKeyable`](#preventing-memoization). +::: ## How data fingerprinting works -For each data value (function argument, `deps` value, or context value), CocoIndex derives a canonical form with this precedence: +For each data value, CocoIndex derives a canonical form with this precedence: -1. If the object implements **`__coco_memo_key__()`**, CocoIndex uses its return value. -2. Otherwise, if you registered a **memo key function** for the object's type, CocoIndex uses that. -3. Otherwise, CocoIndex falls back to structural canonicalization for a limited set of primitives/containers. +1. If the object implements **`__coco_memo_key__()`**, CocoIndex uses its return value. This applies to class instances, not class objects (`type`). +2. Otherwise, if a **memo key function** is registered for the object's type, CocoIndex uses it. +3. Otherwise, CocoIndex uses its built-in canonicalization. The following types are handled automatically (no custom key needed): @@ -24,7 +44,7 @@ The following types are handled automatically (no custom key needed): - **Containers**: `list`, `tuple`, `dict`, `set`, `frozenset` (recursively canonicalized) - **Dataclass instances**: all fields included in definition order - **Pydantic v2 models**: all fields included -- **Class objects** (`type`): identified by module and qualified name +- **Class objects** (`type`): identified by a stable type ID when configured, otherwise by module and qualified name - **Other picklable objects**: used as a fallback via `pickle` The canonical forms are combined into a deterministic fingerprint. If the fingerprint matches a cached entry, the cached result is reused — unless **memo states** indicate it's stale (see [Memo state validation](#memo-state-validation) below). @@ -59,40 +79,126 @@ class UserRow: return ("users", self.user_id, self.updated_at) ``` +The method is used whenever a `UserRow` is fingerprinted as a function argument, `deps` value, or change-detected context value. For example: + +```python +import cocoindex as coco + +@coco.fn(memo=True) +def format_user(row: UserRow) -> str: + print("format_user executed") + return f"{row.user_id}: updated at {row.updated_at}" +``` + +Inside a processing component, the first update executes `format_user`. When logic and context also remain unchanged, a later update with another `UserRow` containing the same `user_id` and `updated_at` reuses the cached result. Changing either value changes the memo key, so the function executes again. + ### Register a key function (when you don't control the type) -If you can't add `__coco_memo_key__` (stdlib / third-party types), register a handler: +If you can't add `__coco_memo_key__` (stdlib / third-party types), register a key function: ```python from pathlib import Path -from cocoindex import register_memo_key_function +from cocoindex import register_memo_type def path_key(p: Path) -> object: p = p.resolve() st = p.stat() return (str(p), st.st_mtime_ns, st.st_size) -register_memo_key_function(Path, path_key) +register_memo_type(Path, path_key) ``` - Registration is **MRO-aware**: if you register both a base class and a subclass, the **most specific** match wins. - Your key function must return the same kinds of stable objects as `__coco_memo_key__` (small primitives/tuples). +### Use stable type IDs across refactors + +By default, CocoIndex identifies types in memo keys by their module and qualified name (`module.QualName`). If you move or rename a class, set `__coco_memo_type_id__` with an ordinary string (for new entries) or `coco.prev_type_id()` (for existing automatic entries): + +```python +from dataclasses import dataclass +from typing import ClassVar +import cocoindex as coco + +# Ordinary string for new entries: +@dataclass +class NewEntry: + __coco_memo_type_id__: ClassVar[str] = "com.example.NewEntry/v1" + + id: str + content: str + +# Retain existing automatic entries: +@dataclass +class SourceEntry: # Moved to new_package.models + __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id( + "old_package.models", + "SourceEntry", + ) + + name: str + version: int +``` + +Declare `prev_type_id` on each moved class; like any class attribute, a base-class declaration propagates to subclasses. + +For types you don't control, register the ID with `register_memo_type()`: + +```python +from some_library import ExternalRecord +import cocoindex as coco + +coco.register_memo_type( + ExternalRecord, + lambda record: (record.id, record.version), + stable_type_id="com.example.ExternalRecord/v1", +) +``` + +Each call replaces the type's entire registration. If a type needs a key function, state function, and stable type ID, pass them together in one call. + +- **Inheritance**: Declaring `__coco_memo_type_id__` on a base class is a deliberate statement that the family shares the same namespace, and propagates to subclasses. For dataclass and Pydantic instances, field tuples keep payloads distinct; a subclass can and should override the ID by declaring its own `__coco_memo_type_id__` when semantics diverge. +- **Precedence**: Type identity resolution follows this precedence: + 1. **Declared stable type ID** (`__coco_memo_type_id__`, own or inherited). + 2. **Registered stable type ID** (most specific registration across the type's MRO via `register_memo_type()`). + 3. **Automatic fallback** (the type's canonical module and qualified name, `module.QualName`). + +:::note[Pickle fallback] +Types using pickle fallback fingerprint raw bytes without type namespacing. Setting or registering a stable type ID alone has no effect on them — register a memo key function together with the stable type ID. +::: + +### Class objects as inputs + +When passing a class object itself (such as `ProductRow` rather than an instance `ProductRow(...)`), CocoIndex fingerprints only the class identity (its declared/registered stable ID or module + qualname), without invoking `__coco_memo_key__` methods or hashing class attributes. + +To customize memoization for class objects globally, register a memo key function on the class's custom metaclass or on `type`. + +:::caution[Avoid colliding class objects] +Unlike instances (where field tuples keep payloads distinct), raw class objects do not hash schema attributes or field tuples. Classes sharing an inherited or metaclass-level stable type ID will produce identical fingerprints. If a memoized function accepts class objects and must distinguish between them, declare a distinct `__coco_memo_type_id__` on each class or provide a custom `memo_key=` at the call site. +::: + ### Override at the call site with `memo_key=` -The two approaches above are **type-level**: every memoized function sees the same fingerprint for a given object. To customize fingerprinting **only for a specific function**, pass `memo_key=` to `@coco.fn` (or `@coco.fn.as_async`). It maps parameter names to either a callable (transform the value before fingerprinting) or `None` (exclude the parameter entirely). +The approaches above are **type-level**: every memoized function sees the same fingerprint for a given object. To customize fingerprinting **only for a specific function**, pass `memo_key=` to `@coco.fn` (or `@coco.fn.as_async`). It maps parameter names to either a callable (transform the value before fingerprinting) or `None` (exclude the parameter entirely). ```python -@coco.fn(memo=True, memo_key={"entry": lambda e: (e.name, e.version), "extra": None}) -def transform(entry: SourceDataEntry, extra: str) -> str: - ... +@coco.fn( + memo=True, + memo_key={ + "entry": lambda entry: (entry.name, entry.version), + "request_id": None, + }, +) +def transform(entry: SourceDataEntry, request_id: str) -> str: + # request_id is request metadata and does not affect the result. + return f"{entry.name}@{entry.version}" ``` Each entry in `memo_key`: - **Callable** — applied to the argument; its return value is fingerprinted in place of the original. Semantically the same as `__coco_memo_key__()` on the type, but scoped to this one function. Useful when the type's default fingerprint is correct everywhere *else*, and only this function should treat the argument differently. - **`None`** — the parameter is excluded from the memo key. Changing its value never invalidates the cache. Useful for arguments that don't affect the result (logger handles, debug flags, request-scoped context that isn't part of the computation). -- **Not listed** — the parameter is fingerprinted normally (type-level `__coco_memo_key__`, registered handler, or default canonicalization). +- **Not listed** — the parameter is fingerprinted normally (type-level `__coco_memo_key__`, registered memo key function, or default canonicalization). It works for every parameter kind: @@ -103,11 +209,6 @@ It works for every parameter kind: `memo_key` is validated at decoration time: unknown parameter names raise `ValueError`, and values that are neither callable nor `None` raise `TypeError`. -:::tip[Picking the right tool] -- You control the type and the fingerprint is the same wherever it's used → **`__coco_memo_key__`**. -- You don't control the type but want a global handler → **`register_memo_key_function`**. -- Only *this* function should treat an argument specially (transform or skip it) → **`memo_key=`**. -::: ## Memo state validation @@ -117,14 +218,14 @@ Sometimes fingerprint matching alone isn't enough to decide whether a cached res - **Async validation**: for an S3 object, send a HEAD request to check freshness — an inherently async operation. - **Stateful validation**: for HTTP resources, store the last fetch time and use `If-Modified-Since` on the next run. -Memo state validation addresses these by letting you attach a **state function** to your objects. It runs *after* a fingerprint match, giving you a chance to check freshness before the cached result is reused. +Memo state validation addresses these by letting you attach a **state function** to your objects. A state function captures initial state when a memo entry is created and, after a later fingerprint match, checks freshness before the cached result is reused. For ordinary function arguments, initial state is captured after a cache miss executes the function. For change-detected context values, initial state is captured when the value is provided. ### How it works -When CocoIndex finds a fingerprint match, it calls each state function with the stored state from the previous run: +CocoIndex calls each state function with the state from the previous invocation: -1. **First run** (no previous state): `prev_state` is `coco.NON_EXISTENCE`. Use `coco.is_non_existence(prev_state)` to detect this. -2. **Subsequent runs**: `prev_state` is whatever you returned last time. +1. **Initial state capture**: `prev_state` is `coco.NON_EXISTENCE`. Use `coco.is_non_existence(prev_state)` to detect this. +2. **Validation after a later fingerprint match**: `prev_state` is whatever the state function returned previously. Your state function returns a `coco.MemoStateOutcome(state=..., memo_valid=...)`: @@ -159,7 +260,9 @@ class LocalFile: # Identity only — which file is it? return str(self.path.resolve()) - def __coco_memo_state__(self, prev_state: tuple[int, str] | coco.NonExistenceType) -> coco.MemoStateOutcome: + def __coco_memo_state__( + self, prev_state: tuple[int, str] | coco.NonExistenceType + ) -> coco.MemoStateOutcome: st = os.stat(self.path) new_mtime = st.st_mtime_ns if coco.is_non_existence(prev_state): @@ -173,8 +276,10 @@ class LocalFile: # mtime unchanged — definitely reusable, no content read needed return coco.MemoStateOutcome(state=prev_state, memo_valid=True) # mtime changed — read content and check hash - content_hash = coco.connectorkits.fingerprint_bytes(self.path.read_bytes()) - return coco.MemoStateOutcome(state=(new_mtime, content_hash), memo_valid=content_hash == prev_hash) + content_hash = hashlib.sha256(self.path.read_bytes()).hexdigest() + return coco.MemoStateOutcome( + state=(new_mtime, content_hash), memo_valid=content_hash == prev_hash + ) ``` :::tip[Keys vs states for files] @@ -189,30 +294,42 @@ This works for simple cases. State validation becomes useful when you need multi ### Register a state function (when you don't control the type) -Pass a `state_fn` keyword argument to `register_memo_key_function`. The state function receives the object as its first argument and `prev_state` as its second. Annotate `prev_state` with the expected type: +Pass a `state_fn` keyword argument together with a callable `key_fn` to `register_memo_type`; `state_fn` cannot be registered by itself. The state function receives the object as its first argument and `prev_state` as its second. Annotate `prev_state` with the expected type: ```python from pathlib import Path -from cocoindex import register_memo_key_function +import cocoindex as coco def path_key(p: Path) -> object: return str(p.resolve()) -def path_state(p: Path, prev_state: tuple[int, int] | coco.NonExistenceType) -> coco.MemoStateOutcome: +def path_state( + p: Path, prev_state: tuple[int, int] | coco.NonExistenceType +) -> coco.MemoStateOutcome: st = p.stat() new_state = (st.st_mtime_ns, st.st_size) memo_valid = not coco.is_non_existence(prev_state) and new_state == prev_state return coco.MemoStateOutcome(state=new_state, memo_valid=memo_valid) -register_memo_key_function(Path, path_key, state_fn=path_state) +coco.register_memo_type( + Path, + path_key, + state_fn=path_state, + stable_type_id="stdlib.pathlib.Path/v1", +) ``` +A registered `state_fn` cannot be layered onto a type's own `__coco_memo_key__`. The intrinsic method takes precedence over registry functions, so neither the registered key function nor its state function is used. If you control the type, define `__coco_memo_state__` alongside `__coco_memo_key__`. The registry has no independent state-only attachment for a third-party type that already defines `__coco_memo_key__`. + +Omit `stable_type_id` if the type does not need a stable type ID. When a type needs both state validation and a stable type ID, provide them in the same registration call. + ### Async state methods A state method can return an `Awaitable`. CocoIndex handles this automatically: - **In an async CocoIndex function**: awaitables from all state methods are gathered concurrently. - **In a sync CocoIndex function**: if no event loop is running, CocoIndex uses `asyncio.run()`. If a loop is already running, it raises an error — switch to an async function or use `@coco.fn.as_async`. +- **For change-detected context values**: initial state is captured synchronously at `builder.provide()`. If `builder.provide()` runs inside an active event loop, use a synchronous state function; an async state function must be initialized by providing the value outside that loop. ```python import cocoindex as coco @@ -252,21 +369,13 @@ class MyStatefulGenerator(coco.NotMemoKeyable): return self._counter ``` -### Register as not memo-keyable (when you don't control the type) - -```python -import cocoindex as coco -from some_library import StatefulGenerator - -coco.register_not_memo_keyable(StatefulGenerator) -``` - -In either case, attempting to use the type as a memo key raises a clear error. +Attempting to use a `NotMemoKeyable` instance as a memo key raises a clear error. ## Best practices - **Keep keys small and deterministic**: use identifiers and versions, not full payloads. No `id(obj)`, pointer addresses, or random values. -- **Separate identity from freshness**: put stable identifiers (file path, URL, primary key) in the key. Put freshness checks (mtime, ETag, version) in the state. +- **Separate identity from freshness**: put stable key values (file path, URL, primary key) in the key. Put freshness checks (mtime, ETag, version) in the state. - **Use state validation for expensive checks**: if freshness validation is costly (content hashing, network calls), a state function lets you do it only when the fingerprint matches, and only when a cheap pre-check (mtime) fails. - **Use `MemoStateOutcome(state=new_state, memo_valid=True)` for cheap state updates**: when a cheap property changes (mtime) but the expensive check (content hash) confirms nothing meaningful changed, return `memo_valid=True` while updating the state. This avoids re-executing the function and avoids re-checking the expensive property next time. +- **Choose stable type IDs deliberately**: use a globally unique, durable name, and bump its version when two type/schema versions should not share memo entries. Do not reuse a stable type ID for unrelated types. - **Mark stateful types as `NotMemoKeyable`**: prevent subtle bugs from incorrect memoization of types with side effects. diff --git a/docs/src/content/docs/programming_guide/function.mdx b/docs/src/content/docs/programming_guide/function.mdx index 749e276ff..b53a48cb4 100644 --- a/docs/src/content/docs/programming_guide/function.mdx +++ b/docs/src/content/docs/programming_guide/function.mdx @@ -77,6 +77,8 @@ A memoized function: When a memoized function's cache hits, its **body does not run** — and neither do any nested `@coco.fn` calls inside it. The cached output is replayed directly: the previous return value is returned, and any target states the function declared on its previous run are carried over. +The fragment below omits the enclosing `App` and processing component. Treat the calls as invocations of the same processing component across successive application updates. + ```python @coco.fn(memo=True) async def inner(text: str) -> str: @@ -88,11 +90,11 @@ async def outer(text: str) -> str: print("outer ran") return await inner(text) + "!" -# First call: both print, both bodies execute +# First component update: both print, both bodies execute. result = await outer("hello") -# Second call with same inputs: nothing prints — outer's cached value is returned -# directly, and inner is not invoked at all. +# Later component update with the same input: nothing prints. The cached value +# from outer is returned directly, and inner is not invoked at all. result = await outer("hello") ``` @@ -174,11 +176,19 @@ def summarize(text: str) -> str: return call_llm(SYSTEM_PROMPT, text, model=MODEL) ``` -The value is canonicalized through the [memoization-key pipeline](../advanced_topics/memoization_keys), which honors `__coco_memo_key__()`, registered memo key functions, and the standard handling for primitives, dataclasses, and Pydantic models. +The value is canonicalized through the [memoization-key pipeline](../advanced_topics/memoization_keys), including custom memo keys, registered memo key functions, stable type IDs, and built-in handling for primitives, containers, dataclasses, Pydantic models, and class objects. -:::caution[Snapshotted at decoration time] -`deps` is evaluated **once** when the decorator is applied (typically at module import), not re-evaluated per call. For per-call or per-instance values — instance attributes in a bound method, request-scoped config, anything that changes at runtime — pass them as regular function arguments instead, so the memoization layer observes each new value. -::: +:::::caution[Snapshotted at decoration time] +`deps` is evaluated **once** when the decorator is applied (typically at module import), not re-evaluated per call. + +Only `deps` is snapshotted at decoration time; regular function arguments are fingerprinted when the function is called. + +This has three consequences: + +- Register stable type IDs before decorating any function whose `deps` include values of that type: define `__coco_memo_type_id__` on the class, or call `register_memo_type(..., stable_type_id=...)`, before any `@coco.fn(..., deps=...)` or `@coco.fn.as_async(..., deps=...)` decorator that references those values. +- A later registration can affect future argument fingerprints, but it cannot retroactively update the logic fingerprint already computed for that decorated function. +- For per-call or per-instance values — instance attributes in a bound method, request-scoped config, anything that changes at runtime — pass them as regular function arguments instead, so the memoization layer observes each new value. +::::: `deps` requires `logic_tracking` to be enabled; combining `deps=` with `logic_tracking=None` raises `ValueError`. diff --git a/python/cocoindex/_internal/api.py b/python/cocoindex/_internal/api.py index ca97a3a1a..5ec6a97c8 100644 --- a/python/cocoindex/_internal/api.py +++ b/python/cocoindex/_internal/api.py @@ -103,7 +103,9 @@ from .memo_fingerprint import ( memo_fingerprint, + prev_type_id, register_memo_key_function, + register_memo_type, NotMemoKeyable, ) @@ -943,7 +945,9 @@ class Cursor: "serialize_by_pickle", # .memo_fingerprint "memo_fingerprint", + "prev_type_id", "register_memo_key_function", + "register_memo_type", "NotMemoKeyable", # .pending_marker "MaybePendingS", diff --git a/python/cocoindex/_internal/context_keys.py b/python/cocoindex/_internal/context_keys.py index 725143c8a..77bfacca1 100644 --- a/python/cocoindex/_internal/context_keys.py +++ b/python/cocoindex/_internal/context_keys.py @@ -180,7 +180,7 @@ def provide(self, key: ContextKey[T], value: T) -> T: state_fns: list[StateFnEntry] = [] canonical = _canonicalize( ("context_key", key._key, value), - _seen=None, + state=None, state_methods=state_fns, ) fp = core.fingerprint_simple_object(canonical) diff --git a/python/cocoindex/_internal/function.py b/python/cocoindex/_internal/function.py index ee38aca17..d6e83da3f 100644 --- a/python/cocoindex/_internal/function.py +++ b/python/cocoindex/_internal/function.py @@ -2029,7 +2029,8 @@ def __call__( # type: ignore[misc] string or model identifier. The value is canonicalized via the memoization-key pipeline (see :doc:`/advanced_topics/memoization_keys` for the full contract, - including ``__coco_memo_key__()`` and registered key functions) + including ``__coco_memo_key__()``, registered memo key functions, + stable type IDs, and built-in handling) and folded into the function's logic fingerprint; when the canonical form changes, memoized results are invalidated and the change propagates to callers according to ``logic_tracking`` @@ -2044,6 +2045,13 @@ def __call__( # type: ignore[misc] config, anything that changes at runtime — pass them as regular function arguments instead. + Stable type IDs used by ``deps`` must therefore exist before the + decorator is applied: define ``__coco_memo_type_id__`` on the class + or call ``register_memo_type(..., stable_type_id=...)`` + before decorating. A later registration can affect future argument + fingerprints, but it cannot retroactively update this function's + already-computed logic fingerprint. + Requires ``logic_tracking`` to be enabled; raises ``ValueError`` if combined with ``logic_tracking=None``. @@ -2157,9 +2165,11 @@ def as_async( # type: ignore[misc] "self": Track own code only, not children. None: No function logic tracking (incompatible with ``deps``). deps: Additional value(s) the function logic depends on but that - aren't visible in its body. See :func:`fn` for the full - contract — the value is canonicalized through the memoization - key pipeline and folded into the function's logic fingerprint. + aren't visible in its body. See :func:`fn` for the full contract, + including stable type IDs and their decoration-time registration + order requirement for ``deps``. The value is canonicalized through + the memoization-key pipeline and folded into the function's logic + fingerprint. Batching and runner are fully supported since the result is always async. diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index 07c9c65b6..89cf031c2 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -1,9 +1,9 @@ """ Persistent memoization fingerprinting (implementation). -This module implements the Python-side canonicalization described in -`docs/docs/dev/memo_key.md`, and relies on a single Rust call to hash the final -canonical form into a fixed-size fingerprint. +This module implements Python-side memo-key canonicalization; user-facing +behavior is documented in +`docs/src/content/docs/advanced_topics/memoization_keys.mdx`. """ from __future__ import annotations @@ -26,17 +26,17 @@ ) from .typing import Fingerprintable - _KeyFn = typing.Callable[[typing.Any], typing.Any] _StateFn = typing.Callable[[typing.Any, typing.Any], typing.Any] -class _MemoFns(typing.NamedTuple): - key_fn: _KeyFn +class _MemoTypeRegistry(typing.NamedTuple): + key_fn: _KeyFn | None = None state_fn: _StateFn | None = None + stable_type_id: str | None = None -_memo_fns: dict[type, _MemoFns] = {} +_memo_type_registry: dict[type, _MemoTypeRegistry] = {} class StateFnEntry(typing.NamedTuple): @@ -51,6 +51,24 @@ class StateFnEntry(typing.NamedTuple): call: typing.Callable[[typing.Any], typing.Any] +@dataclasses.dataclass(slots=True) +class _CanonicalizeState: + seen: dict[int, int] = dataclasses.field(default_factory=dict) + keepalive: list[object] = dataclasses.field(default_factory=list) + + def remember(self, obj: object) -> int | None: + oid = id(obj) + ordinal = self.seen.get(oid) + if ordinal is not None: + assert self.keepalive[ordinal] is obj + return ordinal + + ordinal = len(self.keepalive) + self.keepalive.append(obj) + self.seen[oid] = ordinal + return None + + @functools.cache def _make_state_deserialize_fn( raw_state_fn: typing.Callable[..., typing.Any], @@ -71,7 +89,7 @@ def _make_state_deserialize_fn( ann, source_label=f"prev_state param of {fn_label}()", ) - except Exception: + except Exception: # noqa: BLE001 return make_deserialize_fn(typing.Any) @@ -118,6 +136,169 @@ def canonical_module_name(obj: typing.Any) -> str: return mod +class _PreviousTypeId(str): + """A prior automatic type identity carried through the stable-ID path.""" + + __slots__ = ("_identity_parts",) + _identity_parts: tuple[str, str] + + def __new__(cls, module: str, qualname: str) -> typing.Self: + marker = super().__new__(cls, f"{len(module)}:{module}{qualname}") + object.__setattr__(marker, "_identity_parts", (module, qualname)) + return marker + + def __setattr__(self, name: str, value: object) -> typing.NoReturn: + raise AttributeError(f"{type(self).__name__} is immutable") + + def __delattr__(self, name: str) -> typing.NoReturn: + raise AttributeError(f"{type(self).__name__} is immutable") + + def __reduce__(self) -> tuple[type[_PreviousTypeId], tuple[str, str]]: + return type(self), self._identity_parts + + +def prev_type_id(module: str, qualname: str) -> str: + """Return a marker that reuses a type's prior automatic identity.""" + return _PreviousTypeId(module, qualname) + + +def _register_memo_type_registry(typ: type, registry: _MemoTypeRegistry) -> None: + """Register memo configuration for one exact Python type.""" + _memo_type_registry[typ] = registry + + +def _unregister_memo_type_registry(typ: type) -> None: + """Best-effort removal of an exact-type memo registration.""" + _memo_type_registry.pop(typ, None) + + +def _registered_memo_type_registry(typ: type) -> _MemoTypeRegistry | None: + """Return the exact type's registration from the type-keyed table.""" + return _memo_type_registry.get(typ) + + +_MEMO_KEY_ATTR = "__coco_memo_key__" +_MEMO_STATE_ATTR = "__coco_memo_state__" +_CLASS_OBJECT_OWNER_IDENTITY: tuple[Fingerprintable, Fingerprintable] = ( + canonical_module_name(type), + type.__qualname__, +) + + +def _type_identity_parts( + typ: type, + registry: _MemoTypeRegistry | None, + fallback_owner: type | None = None, + *, + is_class_object: bool = False, +) -> tuple[Fingerprintable, Fingerprintable]: + """Return stable type ID or module+qualname type identity parts. + The stable type ID case still returns two parts to preserve the existing + module/qualname identity shape used by type-aware canonical forms. The + tagged first slot keeps stable type IDs disjoint from ordinary module + names; ``None`` fills the qualname slot. + ``fallback_owner`` anchors the automatic module/qualname identity to the + type that owns a registered key function selected via MRO: the runtime + type's own declared or registered stable ID still wins, but when no stable + ID applies the identity stays with the key function's owner type. + """ + stable_type_id = getattr(typ, "__coco_memo_type_id__", None) + if not isinstance(stable_type_id, str): + if registry is not None and registry.stable_type_id is not None: + stable_type_id = registry.stable_type_id + else: + for owner in typ.__mro__: + if is_class_object and owner is object and typ is not object: + break + reg = ( + registry + if (registry is not None and owner is typ) + else _registered_memo_type_registry(owner) + ) + if reg is not None and reg.stable_type_id is not None: + stable_type_id = reg.stable_type_id + break + if isinstance(stable_type_id, _PreviousTypeId): + return stable_type_id._identity_parts + if isinstance(stable_type_id, str): + return (("__coco_memo_type_id__", stable_type_id), None) + fallback_owner = typ if fallback_owner is None else fallback_owner + return ( + canonical_module_name(fallback_owner), + getattr(fallback_owner, "__qualname__", None), + ) + + +def _canonicalize_registered_memo_key( + obj: object, + owner: type, + registry: _MemoTypeRegistry, + state: _CanonicalizeState, + state_methods: list[StateFnEntry], +) -> Fingerprintable: + key_fn = registry.key_fn + assert key_fn is not None + key = key_fn(obj) + tag = "hook" + if registry.state_fn is not None: + tag = "shook" + bound = functools.partial(registry.state_fn, obj) + state_methods.append(_make_state_fn_entry(bound, registry.state_fn)) + if isinstance(obj, type): + # Class objects keep the identity of the metaclass owner whose + # registration supplied the key function. + identity = _type_identity_parts(owner, registry, is_class_object=True) + else: + # Instances resolve identity on the runtime type, so a subclass can + # override an inherited registered stable ID by declaring or + # registering its own; the automatic identity falls back to the owner. + typ = type(obj) + identity = _type_identity_parts( + typ, + registry if typ is owner else _registered_memo_type_registry(typ), + fallback_owner=owner, + ) + return ( + tag, + *identity, + _canonicalize(key, state, state_methods), + ) + + +def _canonicalize_class_object( + cls: type, + state: _CanonicalizeState, + state_methods: list[StateFnEntry], +) -> Fingerprintable: + """Canonicalize a class object without invoking memo attributes on it.""" + + metaclass: type = type(cls) + metaclass_registry = _registered_memo_type_registry(metaclass) + for owner in metaclass.__mro__: + if owner is object: + break + registry = ( + metaclass_registry + if owner is metaclass + else _registered_memo_type_registry(owner) + ) + if registry is not None and registry.key_fn is not None: + return _canonicalize_registered_memo_key( + cls, owner, registry, state, state_methods + ) + + cls_registry = ( + metaclass_registry if cls is metaclass else _registered_memo_type_registry(cls) + ) + return ( + "hook", + *_CLASS_OBJECT_OWNER_IDENTITY, + # This synthesized identity is already canonical; do not re-enter memo-key + # dispatch, where a registration on ``object`` could intercept it. + ("seq", _type_identity_parts(cls, cls_registry, is_class_object=True)), + ) + + def _is_dataclass_instance(obj: object) -> bool: """Check if obj is a dataclass instance (not a class).""" return dataclasses.is_dataclass(obj) and not isinstance(obj, type) @@ -130,7 +311,8 @@ def _is_pydantic_model(obj: object) -> bool: def _canonicalize_dataclass( obj: object, - _seen: dict[int, int], + registry: _MemoTypeRegistry | None, + state: _CanonicalizeState, state_methods: list[StateFnEntry], ) -> Fingerprintable: """Canonicalize a dataclass instance. @@ -142,10 +324,9 @@ def _canonicalize_dataclass( fields = dataclasses.fields(obj) # type: ignore[arg-type] return ( "dataclass", - canonical_module_name(typ), - typ.__qualname__, + *_type_identity_parts(typ, registry), tuple( - (field.name, _canonicalize(getattr(obj, field.name), _seen, state_methods)) + (field.name, _canonicalize(getattr(obj, field.name), state, state_methods)) for field in fields ), ) @@ -153,7 +334,8 @@ def _canonicalize_dataclass( def _canonicalize_pydantic( obj: object, - _seen: dict[int, int], + registry: _MemoTypeRegistry | None, + state: _CanonicalizeState, state_methods: list[StateFnEntry], ) -> Fingerprintable: """Canonicalize a Pydantic v2 model instance. @@ -165,10 +347,9 @@ def _canonicalize_pydantic( field_names = obj.__pydantic_fields__.keys() # type: ignore[attr-defined] return ( "pydantic", - canonical_module_name(typ), - typ.__qualname__, + *_type_identity_parts(typ, registry), tuple( - (name, _canonicalize(getattr(obj, name), _seen, state_methods)) + (name, _canonicalize(getattr(obj, name), state, state_methods)) for name in field_names ), ) @@ -193,46 +374,141 @@ def __coco_memo_key__(self) -> typing.NoReturn: ) -def register_memo_key_function( - typ: type, key_fn: _KeyFn, *, state_fn: _StateFn | None = None -) -> None: - """Register a memo key function for a type. +@typing.overload +def register_memo_type( + typ: type, + key_fn: _KeyFn, + *, + state_fn: _StateFn | None = None, + stable_type_id: str | None = None, +) -> None: ... + + +@typing.overload +def register_memo_type( + typ: type, + key_fn: None = None, + *, + state_fn: None = None, + stable_type_id: str, +) -> None: ... - Resolution is MRO-aware: the most specific registered base type wins. - If *state_fn* is provided it is stored separately and used for memo state - validation (see ``_canonicalize``). +def register_memo_type( + typ: type, + key_fn: _KeyFn | None = None, + *, + state_fn: _StateFn | None = None, + stable_type_id: str | None = None, +) -> None: + """Register a memo key function, state function, and/or stable type ID for a type. + + Key-function and stable-type-ID resolutions are MRO-aware: the most specific + registered base type wins. Each call replaces the full registration for + ``typ``: omitting ``stable_type_id`` clears any previous registered stable + type ID, and omitting ``key_fn`` or passing ``None`` clears any previous + key/state functions. + + When a registered stable type ID should affect a value used in ``deps=``, + call this before the corresponding ``@coco.fn`` / ``@coco.fn.as_async`` + decorator is applied because ``deps`` fingerprints are computed at + decoration time. """ - _memo_fns[typ] = _MemoFns(key_fn, state_fn) + if not isinstance(typ, type): + raise TypeError( + f"register_memo_type() expects typ to be a type, got {type(typ).__name__}" + ) + if key_fn is None: + if state_fn is not None: + raise TypeError( + "register_memo_type() state_fn requires a memo key function" + ) + if stable_type_id is None: + raise TypeError("register_memo_type() requires a key_fn or stable_type_id") + elif not callable(key_fn): + raise TypeError( + f"register_memo_type() key_fn must be callable, got {type(key_fn).__name__}" + ) + if state_fn is not None and not callable(state_fn): + raise TypeError( + "register_memo_type() state_fn must be callable, " + f"got {type(state_fn).__name__}" + ) + if stable_type_id is not None and not isinstance(stable_type_id, str): + raise TypeError( + "register_memo_type() stable_type_id must be a str, " + f"got {type(stable_type_id).__name__}" + ) + _register_memo_type_registry( + typ, + _MemoTypeRegistry( + key_fn=key_fn, + state_fn=state_fn, + stable_type_id=stable_type_id, + ), + ) -def register_not_memo_keyable(typ: type) -> None: - """Register a type as not memo-keyable. - Use this for third-party types that maintain internal state incompatible - with memoization, but which you cannot modify to inherit from `NotMemoKeyable`. +def register_memo_key_function( + typ: type, + key_fn: _KeyFn, + *, + state_fn: _StateFn | None = None, +) -> None: + """Register a memo key function for a type, preserving any registered stable type ID.""" + existing = _registered_memo_type_registry(typ) if isinstance(typ, type) else None + stable_type_id = existing.stable_type_id if existing is not None else None + register_memo_type( + typ, key_fn=key_fn, state_fn=state_fn, stable_type_id=stable_type_id + ) - Example: - import cocoindex as coco - from some_library import StatefulGenerator - coco.register_not_memo_keyable(StatefulGenerator) +def register_not_memo_keyable(typ: type) -> None: + """Register a type as not memo-keyable. + + Internal helper for tests and internal registrations. It is intentionally + not exported through the public ``cocoindex`` namespace until registered + not-memo-keyable precedence is fixed for types that define + ``__coco_memo_key__`` or otherwise supply memo-key behavior. Public code + should inherit from ``coco.NotMemoKeyable`` when the type is user-owned. """ + if not isinstance(typ, type): + raise TypeError( + "register_not_memo_keyable() expects typ to be a type, " + f"got {type(typ).__name__}" + ) + def _raise_not_memo_keyable(obj: object) -> typing.NoReturn: raise TypeError( f"{type(obj).__name__} cannot be used as a memoization key. " "This type maintains internal state that is incompatible with memoization." ) - _memo_fns[typ] = _MemoFns(_raise_not_memo_keyable) + _register_memo_type_registry(typ, _MemoTypeRegistry(key_fn=_raise_not_memo_keyable)) + + +def unregister_memo_type(typ: type) -> None: + """Remove registered memo configuration for a type.""" + + if isinstance(typ, type): + _unregister_memo_type_registry(typ) def unregister_memo_key_function(typ: type) -> None: - """Remove a previously registered memo key function (best-effort).""" + """Remove a registered memo key function, preserving any registered stable type ID.""" - _memo_fns.pop(typ, None) + if not isinstance(typ, type): + return + reg = _registered_memo_type_registry(typ) + if reg is None: + return + if reg.stable_type_id is not None: + register_memo_type(typ, stable_type_id=reg.stable_type_id) + else: + _unregister_memo_type_registry(typ) def _stable_sort_key(v: Fingerprintable) -> tuple[typing.Any, ...]: @@ -267,12 +543,11 @@ def _stable_sort_key(v: Fingerprintable) -> tuple[typing.Any, ...]: def _canonicalize( obj: object, - _seen: dict[int, int] | None, + state: _CanonicalizeState | None, state_methods: list[StateFnEntry], ) -> Fingerprintable: - # 0) Cycle / shared-reference tracking for containers - if _seen is None: - _seen = {} + if state is None: + state = _CanonicalizeState() # 1) Primitives if obj is None: @@ -283,86 +558,81 @@ def _canonicalize( if isinstance(obj, (bytearray, memoryview)): return bytes(obj) - # 2) Hook / registry (apply once, then recurse on returned key fragment) - hook = getattr(obj, "__coco_memo_key__", None) + # 2) Memo key dispatch. Raw class objects skip memo attributes but honor + # explicit key registrations on their metaclass MRO. + if isinstance(obj, type): + return _canonicalize_class_object(obj, state, state_methods) + + typ = type(obj) + hook = getattr(obj, _MEMO_KEY_ATTR, None) + registry = _registered_memo_type_registry(typ) if hook is not None and callable(hook): k = hook() - typ = type(obj) tag = "hook" - state_hook = getattr(obj, "__coco_memo_state__", None) + state_hook = getattr(obj, _MEMO_STATE_ATTR, None) if state_hook is not None and callable(state_hook): tag = "shook" # raw function for type hint extraction (unbound method on class) - raw_fn = getattr(typ, "__coco_memo_state__") + raw_fn = getattr(typ, _MEMO_STATE_ATTR) state_methods.append(_make_state_fn_entry(state_hook, raw_fn)) return ( tag, - canonical_module_name(typ), - typ.__qualname__, - _canonicalize(k, _seen, state_methods), + *_type_identity_parts(typ, registry), + _canonicalize(k, state, state_methods), ) - for base in type(obj).__mro__: - memo = _memo_fns.get(base) - if memo is not None: - k = memo.key_fn(obj) - tag = "hook" - if memo.state_fn is not None: - tag = "shook" - bound = functools.partial(memo.state_fn, obj) - state_methods.append(_make_state_fn_entry(bound, memo.state_fn)) - return ( - tag, - canonical_module_name(base), - base.__qualname__, - _canonicalize(k, _seen, state_methods), + for owner in typ.__mro__: + owner_registry = ( + registry if owner is typ else _registered_memo_type_registry(owner) + ) + if owner_registry is not None and owner_registry.key_fn is not None: + return _canonicalize_registered_memo_key( + obj, owner, owner_registry, state, state_methods ) # 3) Cycle / shared-reference tracking # # Note: we intentionally do this before branching on container types, so the # logic is shared and we support cyclic/self-referential structures. - oid = id(obj) - ordinal = _seen.get(oid) + ordinal = state.remember(obj) if ordinal is not None: return ("ref", ordinal) - _seen[oid] = len(_seen) # 4) Containers if isinstance(obj, typing.Sequence): - return ("seq", tuple(_canonicalize(e, _seen, state_methods) for e in obj)) + return ("seq", tuple(_canonicalize(e, state, state_methods) for e in obj)) if isinstance(obj, typing.Mapping): items: list[tuple[Fingerprintable, Fingerprintable]] = [] for k, v in obj.items(): items.append( ( - _canonicalize(k, _seen, state_methods), - _canonicalize(v, _seen, state_methods), + _canonicalize(k, state, state_methods), + _canonicalize(v, state, state_methods), ) ) items.sort(key=lambda kv: (_stable_sort_key(kv[0]), _stable_sort_key(kv[1]))) return ("map", tuple(items)) if isinstance(obj, (set, frozenset)): - elts = [_canonicalize(e, _seen, state_methods) for e in obj] + elts = [_canonicalize(e, state, state_methods) for e in obj] elts.sort(key=_stable_sort_key) return ("set", tuple(elts)) # 5) Dataclass instances if _is_dataclass_instance(obj): - return _canonicalize_dataclass(obj, _seen, state_methods) + return _canonicalize_dataclass(obj, registry, state, state_methods) # 6) Pydantic v2 models if _is_pydantic_model(obj): - return _canonicalize_pydantic(obj, _seen, state_methods) + return _canonicalize_pydantic(obj, registry, state, state_methods) # 7) Fallback try: payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) # Tag to avoid colliding with user-provided raw bytes. return ("pickle", payload) - except Exception: + except Exception: # noqa: BLE001 raise TypeError( f"Unsupported type for memoization key: {type(obj)!r}. " "Provide __coco_memo_key__() or register a memo key function." @@ -383,13 +653,13 @@ def _make_call_canonical( getattr(func, "__qualname__", None), ) canonical_args = tuple( - _canonicalize(a, _seen=None, state_methods=state_methods) for a in prefix_args + _canonicalize(a, state=None, state_methods=state_methods) for a in prefix_args ) canonical_args = canonical_args + tuple( - _canonicalize(a, _seen=None, state_methods=state_methods) for a in args + _canonicalize(a, state=None, state_methods=state_methods) for a in args ) canonical_kwargs = tuple( - (k, _canonicalize(v, _seen=None, state_methods=state_methods)) + (k, _canonicalize(v, state=None, state_methods=state_methods)) for k, v in sorted(kwargs.items()) ) return ( @@ -405,7 +675,7 @@ def memo_fingerprint(obj: object) -> core.Fingerprint: # State methods are meaningless for an object-only fingerprint; collect # into a throwaway list so the canonicalizer signature stays uniform. return core.fingerprint_simple_object( - _canonicalize(obj, _seen=None, state_methods=[]) + _canonicalize(obj, state=None, state_methods=[]) ) @@ -440,18 +710,14 @@ def fingerprint_call( return core.fingerprint_simple_object(call_key_obj) -# Register memo key for class types. -register_memo_key_function( - type, - lambda cls: (canonical_module_name(cls), getattr(cls, "__qualname__", None)), -) - - __all__ = [ "NotMemoKeyable", + "fingerprint_call", + "memo_fingerprint", + "prev_type_id", "register_memo_key_function", + "register_memo_type", "register_not_memo_keyable", "unregister_memo_key_function", - "fingerprint_call", - "memo_fingerprint", + "unregister_memo_type", ] diff --git a/python/tests/core/test_function_memo.py b/python/tests/core/test_function_memo.py index 4a5b73b8d..fb7eeec6b 100644 --- a/python/tests/core/test_function_memo.py +++ b/python/tests/core/test_function_memo.py @@ -27,6 +27,43 @@ def __coco_memo_key__(self) -> object: return (self.name, self.version) +@dataclass(frozen=True) +class _StableOldEntry: + name: str + version: int + content: str + + __coco_memo_type_id__ = "test.FunctionMemoEntry/v1" + + +@dataclass(frozen=True) +class _StableNewEntry: + name: str + version: int + content: str + + __coco_memo_type_id__ = "test.FunctionMemoEntry/v1" + + +@dataclass(frozen=True) +class _PreviousIdentityOldEntry: + name: str + version: int + content: str + + +@dataclass(frozen=True) +class _PreviousIdentityMovedEntry: + name: str + version: int + content: str + + __coco_memo_type_id__ = coco.prev_type_id( + _PreviousIdentityOldEntry.__module__, + _PreviousIdentityOldEntry.__qualname__, + ) + + @dataclass class DictSourceDataEntry: name: str @@ -38,6 +75,8 @@ def __coco_memo_key__(self) -> object: _plain_source_data: dict[str, SourceDataEntry] = {} +_stable_type_source_data: dict[str, object] = {} + _dict_source_data: dict[str, DictSourceDataEntry] = {} _metrics = Metrics() @@ -55,6 +94,89 @@ def _process_plain_source_data() -> None: coco.declare_target_state(GlobalDictTarget.target_state(key, transformed_value)) +@coco.fn(memo=True) +def _transform_stable_type_entry(entry: object) -> str: + _metrics.increment("call.transform_stable_type_entry") + return f"processed: {getattr(entry, 'content')}" + + +@coco.fn +def _process_stable_type_source_data() -> None: + for key, value in _stable_type_source_data.items(): + transformed_value = _transform_stable_type_entry(value) + coco.declare_target_state(GlobalDictTarget.target_state(key, transformed_value)) + + +def test_stable_type_id_reuses_memo_across_renamed_dataclass() -> None: + GlobalDictTarget.store.clear() + _stable_type_source_data.clear() + _metrics.clear() + + app = coco.App( + coco.AppConfig( + name="test_stable_type_id_reuses_memo_across_renamed_dataclass", + environment=coco_env, + ), + _process_stable_type_source_data, + ) + + _stable_type_source_data["A"] = _StableOldEntry( + name="A", version=1, content="contentA1" + ) + app.update_blocking() + assert _metrics.collect() == {"call.transform_stable_type_entry": 1} + assert GlobalDictTarget.store.data["A"].data == "processed: contentA1" + + _stable_type_source_data["A"] = _StableNewEntry( + name="A", version=1, content="contentA1" + ) + app.update_blocking() + assert _metrics.collect() == {} + assert GlobalDictTarget.store.data["A"].data == "processed: contentA1" + + _stable_type_source_data["A"] = _StableNewEntry( + name="A", version=2, content="contentA2" + ) + app.update_blocking() + assert _metrics.collect() == {"call.transform_stable_type_entry": 1} + assert GlobalDictTarget.store.data["A"].data == "processed: contentA2" + + +def test_prev_type_id_reuses_existing_memo_after_dataclass_move() -> None: + GlobalDictTarget.store.clear() + _stable_type_source_data.clear() + _metrics.clear() + + app = coco.App( + coco.AppConfig( + name="test_prev_type_id_reuses_existing_memo_after_dataclass_move", + environment=coco_env, + ), + _process_stable_type_source_data, + ) + + _stable_type_source_data["A"] = _PreviousIdentityOldEntry( + name="A", version=1, content="contentA1" + ) + app.update_blocking() + assert _metrics.collect() == {"call.transform_stable_type_entry": 1} + assert GlobalDictTarget.store.data["A"].data == "processed: contentA1" + + _stable_type_source_data["A"] = _PreviousIdentityMovedEntry( + name="A", version=1, content="contentA1" + ) + app.update_blocking() + assert _metrics.collect() == {} + assert GlobalDictTarget.store.data["A"].data == "processed: contentA1" + + _stable_type_source_data["A"] = _PreviousIdentityMovedEntry( + name="A", version=2, content="contentA2" + ) + app.update_blocking() + assert _metrics.collect() == {"call.transform_stable_type_entry": 1} + assert GlobalDictTarget.store.data["A"].data == "processed: contentA2" + + def test_memo_pure_function() -> None: GlobalDictTarget.store.clear() _plain_source_data.clear() diff --git a/python/tests/core/test_memo_state_validation.py b/python/tests/core/test_memo_state_validation.py index aa7df404a..2a2c4544f 100644 --- a/python/tests/core/test_memo_state_validation.py +++ b/python/tests/core/test_memo_state_validation.py @@ -14,6 +14,7 @@ import cocoindex as coco +from cocoindex._internal.memo_fingerprint import unregister_memo_key_function from tests import common from tests.common.target_states import ( DictDataWithPrev, @@ -713,3 +714,158 @@ def test_state_changed_but_reusable_component_sync() -> None: ) app.update_blocking() assert _metrics.collect() == {"call.declare_two_level": 1} + + +# ============================================================================ +# Raw class-object metaclass memo methods are ignored (sync) +# ============================================================================ + +_class_object_source: dict[str, type] = {} +_class_object_prev_states: list[Any] = [] + + +class _ClassObjectMemoMeta(type): + def __coco_memo_key__(cls) -> object: + raise AssertionError("class objects must not call metaclass memo key") + + def __coco_memo_state__(cls, prev_state: Any) -> coco.MemoStateOutcome: + _class_object_prev_states.append(prev_state) + raise AssertionError("class objects must not call metaclass memo state") + + +def _make_class_object( + stable_key: str, + state_value: int, + content: str, + stable_type_id: str | None = None, +) -> type: + attrs: dict[str, object] = { + "stable_key": stable_key, + "state_value": state_value, + "content": content, + } + if stable_type_id is not None: + attrs["__coco_memo_type_id__"] = stable_type_id + return _ClassObjectMemoMeta( + f"ClassObject_{stable_key}_{state_value}_{content}", + (), + attrs, + ) + + +@coco.fn(memo=True) +def _transform_class_object(cls: type) -> str: + _metrics.increment("call.transform_class_object") + return f"class content: {getattr(cls, 'content')}" + + +@coco.fn +def _process_class_objects() -> None: + for key, cls in _class_object_source.items(): + result = _transform_class_object(cls) + coco.declare_target_state(GlobalDictTarget.target_state(key, result)) + + +def test_class_object_metaclass_memo_methods_are_ignored_sync() -> None: + """Class objects use stable type IDs and never call metaclass memo methods.""" + GlobalDictTarget.store.clear() + _class_object_source.clear() + _class_object_prev_states.clear() + _metrics.clear() + + app = coco.App( + coco.AppConfig( + name="test_class_object_metaclass_memo_methods_are_ignored_sync", + environment=coco_env, + ), + _process_class_objects, + ) + + stable_type_id = "test.ClassObjectMemo/v1" + cls = _make_class_object("stable", 1, "A", stable_type_id) + + # Run 1: the transform executes and no metaclass memo methods are called. + _class_object_source["row"] = cls + app.update_blocking() + assert _metrics.collect() == {"call.transform_class_object": 1} + assert _class_object_prev_states == [] + assert GlobalDictTarget.store.data == { + "row": DictDataWithPrev( + data="class content: A", prev=[], prev_may_be_missing=True + ) + } + + # Run 2: a distinct class object with the same stable type ID is a cache hit. + # Class schema/state changes are intentionally not part of the class-object + # fingerprint; users who depend on them must include them with `memo_key=`. + _class_object_source["row"] = _make_class_object("renamed", 2, "B", stable_type_id) + app.update_blocking() + assert _metrics.collect() == {} + assert _class_object_prev_states == [] + assert GlobalDictTarget.store.data["row"].data == "class content: A" + + # Run 3: a distinct stable type ID changes the fingerprint, so the transform + # re-executes and the previous target state is made available. + _class_object_source["row"] = _make_class_object( + "changed", 3, "C", "test.ClassObjectMemo/v2" + ) + app.update_blocking() + assert _metrics.collect() == {"call.transform_class_object": 1} + assert _class_object_prev_states == [] + assert GlobalDictTarget.store.data == { + "row": DictDataWithPrev( + data="class content: C", + prev=["class content: A"], + prev_may_be_missing=False, + ) + } + + +def test_class_object_registered_metaclass_memo_state_sync() -> None: + GlobalDictTarget.store.clear() + _class_object_source.clear() + _class_object_prev_states.clear() + _metrics.clear() + + def class_object_key(cls: Any) -> object: + return cls.stable_key + + def class_object_state(cls: Any, prev_state: Any) -> coco.MemoStateOutcome: + return coco.MemoStateOutcome( + state=cls.state_value, + memo_valid=( + not coco.is_non_existence(prev_state) and cls.state_value == prev_state + ), + ) + + coco.register_memo_key_function( + _ClassObjectMemoMeta, + class_object_key, + state_fn=class_object_state, + ) + try: + app = coco.App( + coco.AppConfig( + name="test_class_object_registered_metaclass_memo_state_sync", + environment=coco_env, + ), + _process_class_objects, + ) + + _class_object_source["row"] = _make_class_object("row", 1, "A") + app.update_blocking() + assert _metrics.collect() == {"call.transform_class_object": 1} + assert GlobalDictTarget.store.data["row"].data == "class content: A" + + _class_object_source["row"] = _make_class_object("row", 1, "B") + app.update_blocking() + assert _metrics.collect() == {} + assert GlobalDictTarget.store.data["row"].data == "class content: A" + + _class_object_source["row"] = _make_class_object("row", 2, "C") + app.update_blocking() + assert _metrics.collect() == {"call.transform_class_object": 1} + assert GlobalDictTarget.store.data["row"].data == "class content: C" + assert _class_object_prev_states == [] + finally: + unregister_memo_key_function(_ClassObjectMemoMeta) diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index ab7ca2f32..b7efe53da 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -1,14 +1,21 @@ +import copy import dataclasses import math -from typing import Any +import pickle +import sys +import weakref +from typing import Any, ClassVar, cast import pytest - +from cocoindex._internal import memo_fingerprint as _memo_fingerprint from cocoindex._internal.function import _apply_memo_key, _normalize_memo_key from cocoindex._internal.memo_fingerprint import ( fingerprint_call, register_memo_key_function, + register_memo_type, + register_not_memo_keyable, unregister_memo_key_function, + unregister_memo_type, ) from cocoindex._internal.typing import MemoStateOutcome @@ -81,8 +88,8 @@ def __init__(self, v: object, irrelevant: object) -> None: self.v = v self.irrelevant = irrelevant - register_memo_key_function(Y, lambda y: ("y", y.v)) try: + register_memo_key_function(Y, lambda y: ("y", y.v)) fp_a = fingerprint_call(_dummy_fn, (Y(5, "a"),), {}, []) fp_irrelevant_changed = fingerprint_call(_dummy_fn, (Y(5, "b"),), {}, []) fp_b = fingerprint_call(_dummy_fn, (Y(6, "a"),), {}, []) @@ -121,9 +128,9 @@ class D: def __init__(self, v: object) -> None: self.v = v - register_memo_key_function(C, lambda x: ("same", x.v)) - register_memo_key_function(D, lambda x: ("same", x.v)) try: + register_memo_key_function(C, lambda x: ("same", x.v)) + register_memo_key_function(D, lambda x: ("same", x.v)) assert fingerprint_call(_dummy_fn, (C(1),), {}, []) != fingerprint_call( _dummy_fn, (D(1),), {}, [] ) @@ -132,6 +139,1555 @@ def __init__(self, v: object) -> None: unregister_memo_key_function(D) +def test_dataclass_stable_type_id_reuses_fingerprint_across_module_move() -> None: + def make_entry(module: str) -> type[Any]: + @dataclasses.dataclass + class Entry: + __coco_memo_type_id__ = "test.Entry/v1" + + value: int + + Entry.__module__ = module + return Entry + + OldEntry = make_entry("tests.old_entries") + NewEntry = make_entry("tests.new_entries") + + assert OldEntry.__qualname__ == NewEntry.__qualname__ + assert OldEntry.__module__ != NewEntry.__module__ + + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) != fingerprint_call( + _dummy_fn, (NewEntry(2),), {}, [] + ) + + +def test_prev_type_id_reuses_previous_automatic_identity() -> None: + import cocoindex as coco + + class OldSourceEntry: + def __init__(self, value: int) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("source-entry", self.value) + + class MovedSourceEntry: + __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id( + "old_package.models", "SourceEntry" + ) + + def __init__(self, value: int) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("source-entry", self.value) + + class DifferentModuleSourceEntry: + __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id( + "other_package.models", "SourceEntry" + ) + + def __init__(self, value: int) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("source-entry", self.value) + + class DifferentQualnameSourceEntry: + __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id( + "old_package.models", "OtherSourceEntry" + ) + + def __init__(self, value: int) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("source-entry", self.value) + + class StableStringSourceEntry: + __coco_memo_type_id__: ClassVar[str] = "old_package.models.SourceEntry" + + def __init__(self, value: int) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("source-entry", self.value) + + OldSourceEntry.__module__ = "old_package.models" + OldSourceEntry.__qualname__ = "SourceEntry" + + old_instance_fingerprint = fingerprint_call(_dummy_fn, (OldSourceEntry(1),), {}, []) + moved_instance_fingerprint = fingerprint_call( + _dummy_fn, (MovedSourceEntry(1),), {}, [] + ) + old_class_fingerprint = fingerprint_call(_dummy_fn, (OldSourceEntry,), {}, []) + moved_class_fingerprint = fingerprint_call(_dummy_fn, (MovedSourceEntry,), {}, []) + + assert old_instance_fingerprint == moved_instance_fingerprint + assert old_class_fingerprint == moved_class_fingerprint + assert old_class_fingerprint != old_instance_fingerprint + assert moved_class_fingerprint != moved_instance_fingerprint + + for changed_type in (DifferentModuleSourceEntry, DifferentQualnameSourceEntry): + assert old_instance_fingerprint != fingerprint_call( + _dummy_fn, (changed_type(1),), {}, [] + ) + assert old_class_fingerprint != fingerprint_call( + _dummy_fn, (changed_type,), {}, [] + ) + + assert old_instance_fingerprint != fingerprint_call( + _dummy_fn, (StableStringSourceEntry(1),), {}, [] + ) + assert old_class_fingerprint != fingerprint_call( + _dummy_fn, (StableStringSourceEntry,), {}, [] + ) + + +def test_prev_type_id_uses_canonical_main_module( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import cocoindex as coco + + class OldSourceEntry: + def __coco_memo_key__(self) -> object: + return "source-entry" + + class CanonicalModuleSourceEntry: + __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id("main", "SourceEntry") + + def __coco_memo_key__(self) -> object: + return "source-entry" + + class LiteralMainModuleSourceEntry: + __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id( + "__main__", "SourceEntry" + ) + + def __coco_memo_key__(self) -> object: + return "source-entry" + + main_module = sys.modules["__main__"] + monkeypatch.setattr(main_module, "__file__", "/tmp/main.py", raising=False) + OldSourceEntry.__module__ = "__main__" + OldSourceEntry.__qualname__ = "SourceEntry" + + old_fingerprint = fingerprint_call(_dummy_fn, (OldSourceEntry(),), {}, []) + assert old_fingerprint == fingerprint_call( + _dummy_fn, (CanonicalModuleSourceEntry(),), {}, [] + ) + assert old_fingerprint != fingerprint_call( + _dummy_fn, (LiteralMainModuleSourceEntry(),), {}, [] + ) + + +def test_prev_type_id_marker_is_immutable_and_round_trips() -> None: + import cocoindex as coco + + left = coco.prev_type_id("a.b", "C") + right = coco.prev_type_id("a", "b.C") + assert str(left) != str(right) + assert left != right + + module = "old:package.models" + qualname = "Outer.Source.Entry" + marker = coco.prev_type_id(module, qualname) + assert isinstance(marker, _memo_fingerprint._PreviousTypeId) + + class IdentitySourceEntry: + __coco_memo_type_id__: ClassVar[str] = marker + + assert ( + _memo_fingerprint._type_identity_parts(IdentitySourceEntry, None) + is marker._identity_parts + ) + with pytest.raises(AttributeError, match="immutable"): + marker._identity_parts = ("mutated", "mutated") + with pytest.raises(AttributeError, match="immutable"): + delattr(marker, "_identity_parts") + + for variant in ( + marker, + copy.copy(marker), + copy.deepcopy(marker), + pickle.loads(pickle.dumps(marker)), + ): + assert type(variant) is type(marker) + + class MovedSourceEntry: + __coco_memo_type_id__: ClassVar[str] = variant + + assert _memo_fingerprint._type_identity_parts(MovedSourceEntry, None) == ( + module, + qualname, + ) + + class RegisteredSourceEntry: + pass + + try: + register_memo_type(RegisteredSourceEntry, stable_type_id=marker) + registry = _memo_fingerprint._registered_memo_type_registry( + RegisteredSourceEntry + ) + assert registry is not None + assert _memo_fingerprint._type_identity_parts( + RegisteredSourceEntry, registry + ) == (module, qualname) + finally: + unregister_memo_type(RegisteredSourceEntry) + + ordinary_marker = str(marker) + + class OrdinaryStringSourceEntry: + __coco_memo_type_id__: ClassVar[str] = ordinary_marker + + assert _memo_fingerprint._type_identity_parts(OrdinaryStringSourceEntry, None) == ( + ("__coco_memo_type_id__", ordinary_marker), + None, + ) + + +def test_pydantic_stable_type_id_allows_renamed_model_reuse() -> None: + try: + from pydantic import BaseModel + except ImportError: + pytest.skip("pydantic not installed") + return + + class OldModel(BaseModel): + __coco_memo_type_id__: ClassVar[str] = "test.Model/v1" + + value: int + + class NewModel(BaseModel): + __coco_memo_type_id__: ClassVar[str] = "test.Model/v1" + + value: int + + assert fingerprint_call( + _dummy_fn, (OldModel(value=1),), {}, [] + ) == fingerprint_call(_dummy_fn, (NewModel(value=1),), {}, []) + assert fingerprint_call( + _dummy_fn, (OldModel(value=1),), {}, [] + ) != fingerprint_call(_dummy_fn, (NewModel(value=2),), {}, []) + + +def test_raw_class_object_stable_type_id_never_calls_staticmethod_memo_key() -> None: + class OldEntry: + __coco_memo_type_id__ = "test.RawClass/v1" + + @staticmethod + def __coco_memo_key__() -> object: + raise AssertionError("class-object fingerprint must not call memo key") + + class NewEntry: + __coco_memo_type_id__ = "test.RawClass/v1" + + @staticmethod + def __coco_memo_key__() -> object: + raise AssertionError("class-object fingerprint must not call memo key") + + class ChangedEntry: + __coco_memo_type_id__ = "test.RawClass/v2" + + @staticmethod + def __coco_memo_key__() -> object: + raise AssertionError("class-object fingerprint must not call memo key") + + OldEntry.__module__ = "tests.old_raw_class" + NewEntry.__module__ = "tests.new_raw_class" + ChangedEntry.__module__ = "tests.new_raw_class" + + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry,), {}, [] + ) + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) != fingerprint_call( + _dummy_fn, (ChangedEntry,), {}, [] + ) + + +def test_registered_stable_type_id_applies_to_class_objects() -> None: + class OldEntry: + pass + + class NewEntry: + pass + + class ChangedEntry: + pass + + try: + register_memo_type(OldEntry, stable_type_id="test.RegisteredRawClass/v1") + register_memo_type(NewEntry, stable_type_id="test.RegisteredRawClass/v1") + register_memo_type(ChangedEntry, stable_type_id="test.RegisteredRawClass/v2") + + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry,), {}, [] + ) + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) != fingerprint_call( + _dummy_fn, (ChangedEntry,), {}, [] + ) + finally: + unregister_memo_type(OldEntry) + unregister_memo_type(NewEntry) + unregister_memo_type(ChangedEntry) + + +def test_hook_memo_key_fragment_preserves_parent_cycle() -> None: + class Entry: + def __init__(self, parent: list[object]) -> None: + self.parent = parent + + def __coco_memo_key__(self) -> object: + return self.parent + + def make_graph() -> list[object]: + parent: list[object] = [] + parent.append(Entry(parent)) + return parent + + graph = make_graph() + canonical = _memo_fingerprint._canonicalize(graph, None, []) + + assert canonical == ( + "seq", + (("hook", *_memo_fingerprint._type_identity_parts(Entry, None), ("ref", 0)),), + ) + assert _memo_fingerprint.memo_fingerprint( + graph + ) == _memo_fingerprint.memo_fingerprint(make_graph()) + + +def test_registered_memo_key_fragment_preserves_parent_cycle() -> None: + class Entry: + def __init__(self, parent: list[object]) -> None: + self.parent = parent + + def make_graph() -> list[object]: + parent: list[object] = [] + parent.append(Entry(parent)) + return parent + + try: + register_memo_key_function(Entry, lambda entry: entry.parent) + graph = make_graph() + canonical = _memo_fingerprint._canonicalize(graph, None, []) + + assert canonical == ( + "seq", + ( + ( + "hook", + *_memo_fingerprint._type_identity_parts(Entry, None), + ("ref", 0), + ), + ), + ) + assert _memo_fingerprint.memo_fingerprint( + graph + ) == _memo_fingerprint.memo_fingerprint(make_graph()) + finally: + unregister_memo_key_function(Entry) + + +def test_hook_memo_key_fragments_remain_alive_for_root_traversal() -> None: + class Fragment(list[object]): + pass + + first_fragment_ref: weakref.ReferenceType[Fragment] | None = None + + class FirstEntry: + def __coco_memo_key__(self) -> object: + nonlocal first_fragment_ref + fragment = Fragment(["first"]) + first_fragment_ref = weakref.ref(fragment) + return fragment + + class SecondEntry: + def __coco_memo_key__(self) -> object: + assert first_fragment_ref is not None + assert first_fragment_ref() is not None + return Fragment(["second"]) + + canonical = _memo_fingerprint._canonicalize([FirstEntry(), SecondEntry()], None, []) + + assert canonical == ( + "seq", + ( + ( + "hook", + *_memo_fingerprint._type_identity_parts(FirstEntry, None), + ("seq", ("first",)), + ), + ( + "hook", + *_memo_fingerprint._type_identity_parts(SecondEntry, None), + ("seq", ("second",)), + ), + ), + ) + + +def test_memo_key_fragment_preserves_shared_reference_ordinals() -> None: + class Entry: + def __init__(self) -> None: + self.shared = ["shared"] + + def __coco_memo_key__(self) -> object: + return (self.shared, self.shared) + + top_level = _memo_fingerprint._canonicalize(Entry(), None, []) + parent_wrapped = _memo_fingerprint._canonicalize([Entry()], None, []) + + assert top_level == ( + "hook", + *_memo_fingerprint._type_identity_parts(Entry, None), + ("seq", (("seq", ("shared",)), ("ref", 1))), + ) + assert parent_wrapped == ( + "seq", + ( + ( + "hook", + *_memo_fingerprint._type_identity_parts(Entry, None), + ("seq", (("seq", ("shared",)), ("ref", 2))), + ), + ), + ) + + +def test_raw_class_object_default_identity_never_calls_classmethod_memo_key() -> None: + class Entry: + @classmethod + def __coco_memo_key__(cls) -> object: + raise AssertionError("class-object fingerprint must not call memo key") + + class OtherEntry: + @classmethod + def __coco_memo_key__(cls) -> object: + raise AssertionError("class-object fingerprint must not call memo key") + + Entry.__module__ = "tests.raw_class_default" + OtherEntry.__module__ = "tests.raw_class_default" + + assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == fingerprint_call( + _dummy_fn, (Entry,), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Entry,), {}, []) != fingerprint_call( + _dummy_fn, (OtherEntry,), {}, [] + ) + + +def test_raw_class_object_honors_registered_metaclass_memo_key_and_state() -> None: + key_calls: list[type] = [] + state_calls: list[tuple[type, object]] = [] + + class MemoMeta(type): + def __coco_memo_key__(cls) -> object: + raise AssertionError("raw classes must not call metaclass memo attributes") + + def __coco_memo_state__(cls, prev_state: object) -> MemoStateOutcome: + raise AssertionError("raw classes must not call metaclass memo attributes") + + class OldEntry(metaclass=MemoMeta): + pass + + class NewEntry(metaclass=MemoMeta): + pass + + def metaclass_key(cls: type) -> object: + key_calls.append(cls) + return ("metaclass", cls.__name__) + + def metaclass_state(cls: type, prev_state: object) -> MemoStateOutcome: + state_calls.append((cls, prev_state)) + return MemoStateOutcome( + state=(cls.__name__, prev_state), memo_valid=prev_state == "reusable" + ) + + stable_type_id = "test.RawClassRegisteredMetaOwner/v1" + try: + register_memo_type( + MemoMeta, + metaclass_key, + state_fn=metaclass_state, + stable_type_id=stable_type_id, + ) + old_state_methods: list[Any] = [] + new_state_methods: list[Any] = [] + old_canonical = _memo_fingerprint._canonicalize( + OldEntry, None, old_state_methods + ) + new_canonical = _memo_fingerprint._canonicalize( + NewEntry, None, new_state_methods + ) + + assert old_canonical == ( + "shook", + ("__coco_memo_type_id__", stable_type_id), + None, + ("seq", ("metaclass", OldEntry.__name__)), + ) + assert new_canonical == ( + "shook", + ("__coco_memo_type_id__", stable_type_id), + None, + ("seq", ("metaclass", NewEntry.__name__)), + ) + assert key_calls == [OldEntry, NewEntry] + assert len(old_state_methods) == 1 + assert len(new_state_methods) == 1 + assert old_state_methods[0].call("old previous") == MemoStateOutcome( + state=(OldEntry.__name__, "old previous"), memo_valid=False + ) + assert new_state_methods[0].call("reusable") == MemoStateOutcome( + state=(NewEntry.__name__, "reusable"), memo_valid=True + ) + assert state_calls == [ + (OldEntry, "old previous"), + (NewEntry, "reusable"), + ] + finally: + unregister_memo_type(MemoMeta) + + +def test_prev_type_id_reuses_registered_metaclass_owner_identity() -> None: + import cocoindex as coco + + class OldMemoMeta(type): + pass + + class MovedMemoMeta(type): + pass + + class OldEntry(metaclass=OldMemoMeta): + pass + + class MovedEntry(metaclass=MovedMemoMeta): + pass + + OldMemoMeta.__module__ = "old_package.models" + OldMemoMeta.__qualname__ = "SourceMeta" + + def metaclass_key(cls: type) -> object: + return "source-class" + + try: + register_memo_type(OldMemoMeta, metaclass_key) + register_memo_type( + MovedMemoMeta, + metaclass_key, + stable_type_id=coco.prev_type_id( + "old_package.models", + "SourceMeta", + ), + ) + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( + _dummy_fn, (MovedEntry,), {}, [] + ) + finally: + unregister_memo_type(OldMemoMeta) + unregister_memo_type(MovedMemoMeta) + + +def test_raw_class_object_honors_registered_type_memo_key_and_state() -> None: + key_calls: list[type] = [] + state_calls: list[tuple[type, object]] = [] + + class MemoAttributesMustNotRun: + @classmethod + def __coco_memo_key__(cls) -> object: + raise AssertionError("raw classes must not call class memo attributes") + + @classmethod + def __coco_memo_state__(cls, prev_state: object) -> MemoStateOutcome: + raise AssertionError("raw classes must not call class memo attributes") + + class OldEntry(MemoAttributesMustNotRun): + pass + + class NewEntry(MemoAttributesMustNotRun): + pass + + def type_key(cls: type) -> object: + key_calls.append(cls) + return ("type", cls.__name__) + + def type_state(cls: type, prev_state: object) -> MemoStateOutcome: + state_calls.append((cls, prev_state)) + return MemoStateOutcome( + state=(cls.__name__, prev_state), memo_valid=prev_state == "reusable" + ) + + stable_type_id = "test.RawClassRegisteredTypeOwner/v1" + try: + register_memo_type( + type, + type_key, + state_fn=type_state, + stable_type_id=stable_type_id, + ) + old_state_methods: list[Any] = [] + new_state_methods: list[Any] = [] + old_canonical = _memo_fingerprint._canonicalize( + OldEntry, None, old_state_methods + ) + new_canonical = _memo_fingerprint._canonicalize( + NewEntry, None, new_state_methods + ) + + assert old_canonical == ( + "shook", + ("__coco_memo_type_id__", stable_type_id), + None, + ("seq", ("type", OldEntry.__name__)), + ) + assert new_canonical == ( + "shook", + ("__coco_memo_type_id__", stable_type_id), + None, + ("seq", ("type", NewEntry.__name__)), + ) + assert key_calls == [OldEntry, NewEntry] + assert len(old_state_methods) == 1 + assert len(new_state_methods) == 1 + assert old_state_methods[0].call("old previous") == MemoStateOutcome( + state=(OldEntry.__name__, "old previous"), memo_valid=False + ) + assert new_state_methods[0].call("reusable") == MemoStateOutcome( + state=(NewEntry.__name__, "reusable"), memo_valid=True + ) + assert state_calls == [ + (OldEntry, "old previous"), + (NewEntry, "reusable"), + ] + finally: + unregister_memo_type(type) + + +def test_raw_class_object_ignores_registered_object_memo_key() -> None: + object_key_calls: list[object] = [] + + class Entry: + pass + + class OtherEntry: + pass + + expected_entry = fingerprint_call(_dummy_fn, (Entry,), {}, []) + expected_other = fingerprint_call(_dummy_fn, (OtherEntry,), {}, []) + assert expected_entry != expected_other + + def object_key(obj: object) -> object: + object_key_calls.append(obj) + return "object memo key ran" + + try: + register_memo_type( + object, + object_key, + stable_type_id="test.IgnoredObjectClass/v1", + ) + assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == expected_entry + assert fingerprint_call(_dummy_fn, (OtherEntry,), {}, []) == expected_other + assert object_key_calls == [] + finally: + unregister_memo_type(object) + + +def test_raw_class_object_ignores_registered_type_stable_type_id() -> None: + class Entry: + pass + + Entry.__module__ = "tests.raw_class_registered_type_stable_id" + original = fingerprint_call(_dummy_fn, (Entry,), {}, []) + + try: + register_memo_type(type, stable_type_id="test.RegisteredType/v1") + assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == original + finally: + unregister_memo_type(type) + + +def test_raw_class_object_stable_type_id_is_inherited_by_subclasses() -> None: + class Parent: + __coco_memo_type_id__ = "test.RawClassExact/v1" + + class Child(Parent): + pass + + class OverridingChild(Parent): + __coco_memo_type_id__ = "test.RawClassExactChild/v1" + + assert fingerprint_call(_dummy_fn, (Parent,), {}, []) == fingerprint_call( + _dummy_fn, (Child,), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Parent,), {}, []) != fingerprint_call( + _dummy_fn, (OverridingChild,), {}, [] + ) + + +def test_register_memo_type_registers_key_function_and_stable_type_id_for_owner_base() -> ( + None +): + class Base: + def __init__(self, value: object) -> None: + self.value = value + + class ChildA(Base): + pass + + class ChildB(Base): + pass + + try: + register_memo_type( + Base, + lambda entry: ("base", entry.value), + stable_type_id="test.RegisteredBase/v1", + ) + assert fingerprint_call(_dummy_fn, (ChildA(1),), {}, []) == fingerprint_call( + _dummy_fn, (ChildB(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (ChildA(1),), {}, []) != fingerprint_call( + _dummy_fn, (ChildB(2),), {}, [] + ) + finally: + unregister_memo_type(Base) + + +def test_prev_type_id_reuses_registered_mro_owner_identity() -> None: + import cocoindex as coco + + class OldBase: + def __init__(self, value: object) -> None: + self.value = value + + class OldChild(OldBase): + pass + + class MovedBase: + def __init__(self, value: object) -> None: + self.value = value + + class MovedChild(MovedBase): + pass + + OldBase.__module__ = "old_package.models" + OldBase.__qualname__ = "SourceBase" + + def base_key(entry: OldBase | MovedBase) -> object: + return ("base", entry.value) + + try: + register_memo_type(OldBase, base_key) + register_memo_type( + MovedBase, + base_key, + stable_type_id=coco.prev_type_id( + "old_package.models", + "SourceBase", + ), + ) + assert fingerprint_call(_dummy_fn, (OldChild(1),), {}, []) == fingerprint_call( + _dummy_fn, (MovedChild(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (OldChild(1),), {}, []) != fingerprint_call( + _dummy_fn, (MovedChild(2),), {}, [] + ) + finally: + unregister_memo_type(OldBase) + unregister_memo_type(MovedBase) + + +def test_register_memo_type_registers_stable_type_id_without_key_function() -> None: + class OldEntry: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class NewEntry: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + try: + register_memo_type(OldEntry, stable_type_id="test.RegisteredEntry/v1") + register_memo_type(NewEntry, None, stable_type_id="test.RegisteredEntry/v1") + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) != fingerprint_call( + _dummy_fn, (NewEntry(2),), {}, [] + ) + finally: + unregister_memo_type(OldEntry) + unregister_memo_type(NewEntry) + + +def test_stable_type_id_only_registration_propagates_to_subclasses_across_mro() -> None: + class Parent: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class Child(Parent): + pass + + class OverridingChild(Parent): + pass + + class SameStableTypeId: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + try: + register_memo_type(Parent, stable_type_id="test.RegisteredParent/v1") + register_memo_type( + OverridingChild, stable_type_id="test.RegisteredOverridingChild/v1" + ) + register_memo_type(SameStableTypeId, stable_type_id="test.RegisteredParent/v1") + + # Child inherits Parent's registered stable type ID across MRO: + assert fingerprint_call(_dummy_fn, (Parent(1),), {}, []) == fingerprint_call( + _dummy_fn, (Child(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Parent(1),), {}, []) == fingerprint_call( + _dummy_fn, (SameStableTypeId(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Parent,), {}, []) == fingerprint_call( + _dummy_fn, (Child,), {}, [] + ) + # OverridingChild has its own more specific registration: + assert fingerprint_call(_dummy_fn, (Parent(1),), {}, []) != fingerprint_call( + _dummy_fn, (OverridingChild(1),), {}, [] + ) + finally: + unregister_memo_type(Parent) + unregister_memo_type(OverridingChild) + unregister_memo_type(SameStableTypeId) + + +def test_declared_stable_type_id_is_inherited_and_takes_precedence_over_registered() -> ( + None +): + class Base: + __coco_memo_type_id__: ClassVar[str] = "test.DeclaredBase/v1" + + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class Sub(Base): + pass + + class OverridingSub(Base): + __coco_memo_type_id__: ClassVar[str] = "test.DeclaredOverridingSub/v1" + + try: + # Register an ID on Sub — declared (inherited from Base) must take precedence: + register_memo_type(Sub, stable_type_id="test.RegisteredSub/v1") + + assert _memo_fingerprint._type_identity_parts(Base, None) == ( + ("__coco_memo_type_id__", "test.DeclaredBase/v1"), + None, + ) + # Sub inherits Base's declared ID, beating Sub's registered ID: + assert _memo_fingerprint._type_identity_parts(Sub, None) == ( + ("__coco_memo_type_id__", "test.DeclaredBase/v1"), + None, + ) + assert fingerprint_call(_dummy_fn, (Base(1),), {}, []) == fingerprint_call( + _dummy_fn, (Sub(1),), {}, [] + ) + # OverridingSub overrides Base's declared ID: + assert _memo_fingerprint._type_identity_parts(OverridingSub, None) == ( + ("__coco_memo_type_id__", "test.DeclaredOverridingSub/v1"), + None, + ) + assert fingerprint_call(_dummy_fn, (Base(1),), {}, []) != fingerprint_call( + _dummy_fn, (OverridingSub(1),), {}, [] + ) + finally: + unregister_memo_type(Sub) + + +def test_prev_type_id_on_base_class_propagates_to_subclasses() -> None: + import cocoindex as coco + + class OldBase: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class MovedBase: + __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id( + "old_package.models", "OldBase" + ) + + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class MovedSub(MovedBase): + pass + + OldBase.__module__ = "old_package.models" + OldBase.__qualname__ = "OldBase" + + # End-to-end black-box verification: MovedSub matches OldBase fingerprint: + assert fingerprint_call(_dummy_fn, (MovedSub(1),), {}, []) == fingerprint_call( + _dummy_fn, (OldBase(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (MovedSub(1),), {}, []) != fingerprint_call( + _dummy_fn, (OldBase(2),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (MovedBase(1),), {}, []) == fingerprint_call( + _dummy_fn, (MovedSub(1),), {}, [] + ) + + +def test_stable_type_id_only_subclass_registration_does_not_hide_base_key() -> None: + class Parent: + def __init__(self, value: object, ignored: object) -> None: + self.value = value + self.ignored = ignored + + class Child(Parent): + pass + + try: + register_memo_type(Parent, lambda entry: ("parent", entry.value)) + register_memo_type(Child, stable_type_id="test.RegisteredExactChild/v1") + + assert fingerprint_call(_dummy_fn, (Child(1, "a"),), {}, []) == ( + fingerprint_call(_dummy_fn, (Child(1, "b"),), {}, []) + ) + assert fingerprint_call(_dummy_fn, (Child(1, "a"),), {}, []) != ( + fingerprint_call(_dummy_fn, (Child(2, "a"),), {}, []) + ) + finally: + unregister_memo_type(Child) + unregister_memo_type(Parent) + + +def test_subclass_declared_stable_type_id_overrides_registered_base_identity() -> None: + class Base: + def __init__(self, value: object, ignored: object) -> None: + self.value = value + self.ignored = ignored + + class Sub(Base): + __coco_memo_type_id__: ClassVar[str] = "test.SubDeclaredOverride/v1" + + class TwinSub(Base): + __coco_memo_type_id__: ClassVar[str] = "test.SubDeclaredOverride/v1" + + try: + register_memo_type( + Base, + lambda entry: ("base", entry.value), + stable_type_id="test.BaseRegisteredOverride/v1", + ) + + # Base's registered key function still drives Sub's key shape: + assert fingerprint_call(_dummy_fn, (Sub(1, "a"),), {}, []) == fingerprint_call( + _dummy_fn, (Sub(1, "b"),), {}, [] + ) + # Sub's declared ID overrides the identity inherited from Base's registration: + assert fingerprint_call(_dummy_fn, (Sub(1, "a"),), {}, []) != fingerprint_call( + _dummy_fn, (Base(1, "a"),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Sub(1, "a"),), {}, []) == fingerprint_call( + _dummy_fn, (TwinSub(1, "a"),), {}, [] + ) + finally: + unregister_memo_type(Base) + + +def test_subclass_registered_stable_type_id_overrides_registered_base_identity() -> ( + None +): + class Base: + def __init__(self, value: object, ignored: object) -> None: + self.value = value + self.ignored = ignored + + class Sub(Base): + pass + + class TwinSub(Base): + pass + + try: + register_memo_type( + Base, + lambda entry: ("base", entry.value), + stable_type_id="test.BaseRegisteredOverride/v2", + ) + register_memo_type(Sub, stable_type_id="test.SubRegisteredOverride/v2") + register_memo_type(TwinSub, stable_type_id="test.SubRegisteredOverride/v2") + + # Base's registered key function still drives Sub's key shape: + assert fingerprint_call(_dummy_fn, (Sub(1, "a"),), {}, []) == fingerprint_call( + _dummy_fn, (Sub(1, "b"),), {}, [] + ) + # Sub's registered ID overrides the identity inherited from Base's registration: + assert fingerprint_call(_dummy_fn, (Sub(1, "a"),), {}, []) != fingerprint_call( + _dummy_fn, (Base(1, "a"),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Sub(1, "a"),), {}, []) == fingerprint_call( + _dummy_fn, (TwinSub(1, "a"),), {}, [] + ) + finally: + unregister_memo_type(TwinSub) + unregister_memo_type(Sub) + unregister_memo_type(Base) + + +def test_register_memo_type_replaces_and_clears_omitted_state_fn() -> None: + @dataclasses.dataclass + class Entry: + value: int + ignored: str + + def state_fn(obj: Any, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome(state=("state", obj.value, prev_state), memo_valid=True) + + try: + # 1. Start with (key_fn, state_fn, stable_type_id) + register_memo_type( + Entry, + lambda e: ("key", e.value), + state_fn=state_fn, + stable_type_id="test.Replace/v1", + ) + methods: list[Any] = [] + fp_with_state = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) + assert len(methods) == 1 + assert fp_with_state == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) + + # 2. Re-register omitting state_fn -> clears state_fn while keeping key_fn and stable ID + register_memo_type( + Entry, + lambda e: ("key", e.value), + stable_type_id="test.Replace/v1", + ) + methods = [] + fp_no_state = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) + assert len(methods) == 0 + assert fp_no_state != fp_with_state + assert fp_no_state == fingerprint_call( + _dummy_fn, (Entry(1, "b"),), {}, [] + ) # key_fn still active + + # 3. Re-register with new stable ID -> changes fingerprint while keeping key_fn + register_memo_type( + Entry, + lambda e: ("key", e.value), + stable_type_id="test.Replace/v1_updated", + ) + fp_new_id = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) + assert fp_new_id != fp_no_state + assert fp_new_id == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) + finally: + unregister_memo_type(Entry) + + +def test_register_memo_type_stable_id_only_clears_previous_key_and_state() -> None: + @dataclasses.dataclass + class Entry: + value: int + ignored: str + + @dataclasses.dataclass + class ReferenceSameId: + value: int + ignored: str + + stable_id = "test.Replace/v2" + try: + register_memo_type(ReferenceSameId, stable_type_id=stable_id) + + # 1. Start with key_fn + register_memo_type(Entry, lambda e: ("key", e.value)) + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) + + # 2. Re-register with stable_id only -> clears key_fn (dataclass fields participate again) and sets stable_id + register_memo_type(Entry, stable_type_id=stable_id) + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) != fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) == fingerprint_call(_dummy_fn, (ReferenceSameId(1, "a"),), {}, []) + finally: + unregister_memo_type(Entry) + unregister_memo_type(ReferenceSameId) + + +def test_register_memo_type_key_only_clears_previous_stable_id() -> None: + class Entry: + def __init__(self, value: int) -> None: + self.value = value + + class SameIdEntry: + def __init__(self, value: int) -> None: + self.value = value + + stable_id = "test.Replace/v3" + try: + register_memo_type(SameIdEntry, lambda e: e.value, stable_type_id=stable_id) + + # 1. Start with stable_id + register_memo_type(Entry, lambda e: e.value, stable_type_id=stable_id) + assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) == fingerprint_call( + _dummy_fn, (SameIdEntry(1),), {}, [] + ) + + # 2. Re-register with key_fn only -> clears stable_id (falls back to module.qualname) + register_memo_type(Entry, lambda e: e.value) + assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) != fingerprint_call( + _dummy_fn, (SameIdEntry(1),), {}, [] + ) + finally: + unregister_memo_type(Entry) + unregister_memo_type(SameIdEntry) + + +def test_intrinsic_hooks_and_declared_stable_id_beat_registration() -> None: + declared_type_id = "test.DeclaredIntrinsic/v1" + + class Entry: + __coco_memo_type_id__ = declared_type_id + + def __coco_memo_key__(self) -> object: + return "intrinsic-key" + + def __coco_memo_state__(self, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome(state=prev_state, memo_valid=True) + + def registered_key(_entry: Entry) -> object: + raise AssertionError("intrinsic key must beat the registered key") + + def registered_state(_entry: Entry, _prev_state: object) -> MemoStateOutcome: + raise AssertionError("intrinsic state must beat the registered state") + + try: + register_memo_type( + Entry, + registered_key, + state_fn=registered_state, + stable_type_id="test.RegisteredIntrinsic/v1", + ) + state_methods: list[Any] = [] + canonical = _memo_fingerprint._canonicalize(Entry(), None, state_methods) + assert isinstance(canonical, tuple) + assert canonical[:3] == ( + "shook", + ("__coco_memo_type_id__", declared_type_id), + None, + ) + assert len(state_methods) == 1 + assert state_methods[0].call("previous") == MemoStateOutcome( + state="previous", memo_valid=True + ) + finally: + unregister_memo_type(Entry) + + +def test_declared_stable_id_beats_registration_for_selected_key_owner() -> None: + declared_type_id = "test.DeclaredBaseOwner/v1" + + class BaseEntry: + __coco_memo_type_id__ = declared_type_id + + def __init__(self, value: object) -> None: + self.value = value + + class ChildEntry(BaseEntry): + pass + + try: + register_memo_type( + BaseEntry, + lambda entry: ("registered", entry.value), + stable_type_id="test.RegisteredBaseOwner/v1", + ) + assert _memo_fingerprint._canonicalize(ChildEntry(1), None, []) == ( + "hook", + ("__coco_memo_type_id__", declared_type_id), + None, + ("seq", ("registered", 1)), + ) + finally: + unregister_memo_type(BaseEntry) + + +def test_combined_registration_uses_stable_type_id_and_collects_state_fn() -> None: + class OldEntry: + def __init__(self, value: object) -> None: + self.value = value + + class NewEntry: + def __init__(self, value: object) -> None: + self.value = value + + def key_fn(entry: Any) -> object: + return ("entry", entry.value) + + def state_fn(entry: Any, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome( + state=("state", entry.value, prev_state), memo_valid=True + ) + + try: + register_memo_type( + OldEntry, + key_fn, + state_fn=state_fn, + stable_type_id="test.CombinedStateStable/v1", + ) + register_memo_type( + NewEntry, + key_fn, + state_fn=state_fn, + stable_type_id="test.CombinedStateStable/v1", + ) + + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == ( + fingerprint_call(_dummy_fn, (NewEntry(1),), {}, []) + ) + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) != ( + fingerprint_call(_dummy_fn, (NewEntry(2),), {}, []) + ) + methods: list[Any] = [] + fingerprint_call(_dummy_fn, (OldEntry(1),), {}, methods) + assert len(methods) == 1 + assert methods[0].call("prev").state == ("state", 1, "prev") + finally: + unregister_memo_type(OldEntry) + unregister_memo_type(NewEntry) + + +def test_register_not_memo_keyable_replaces_stable_type_id_for_class_objects() -> None: + class Entry: + pass + + class SameStableTypeId: + pass + + try: + register_memo_type(Entry, stable_type_id="test.NotMemoKeyableReplacesStable/v1") + register_memo_type( + SameStableTypeId, + stable_type_id="test.NotMemoKeyableReplacesStable/v1", + ) + assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == fingerprint_call( + _dummy_fn, (SameStableTypeId,), {}, [] + ) + + register_not_memo_keyable(Entry) + assert fingerprint_call(_dummy_fn, (Entry,), {}, []) != fingerprint_call( + _dummy_fn, (SameStableTypeId,), {}, [] + ) + finally: + unregister_memo_type(Entry) + unregister_memo_type(SameStableTypeId) + + +@pytest.mark.parametrize( + ("args", "kwargs", "match"), + [ + ((), {}, "requires a key_fn or stable_type_id"), + ((object(),), {}, "key_fn must be callable"), + ( + (lambda entry: entry,), + {"state_fn": object()}, + "state_fn must be callable", + ), + ( + (), + {"stable_type_id": 123}, + "stable_type_id must be a str", + ), + ], +) +def test_register_memo_type_rejects_invalid_forms( + args: tuple[Any, ...], + kwargs: dict[str, Any], + match: str, +) -> None: + class Entry: + pass + + with pytest.raises(TypeError, match=match): + register_memo_type(Entry, *args, **kwargs) + + +def test_register_memo_type_rejects_state_fn_without_key_function() -> None: + class Entry: + pass + + def state_fn(obj: Entry, prev_state: object) -> object: + return prev_state + + kwargs: Any = {"state_fn": state_fn} + + with pytest.raises(TypeError, match="state_fn requires a memo key function"): + register_memo_type(Entry, **kwargs) + + +def test_register_and_unregister_memo_key_function_shortcuts_preserve_stable_type_id() -> ( + None +): + @dataclasses.dataclass + class Entry: + value: int + ignored: str + + @dataclasses.dataclass + class Twin: + value: int + ignored: str + + stable_id = "test.ShortcutPreservesStable/v1" + try: + # Register reference twin with stable ID: + register_memo_type(Twin, stable_type_id=stable_id) + + # 1. Register only stable type ID on Entry: + register_memo_type(Entry, stable_type_id=stable_id) + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) != fingerprint_call( + _dummy_fn, (Entry(1, "b"),), {}, [] + ) # No key_fn yet: dataclass hashing includes 'ignored' + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) == fingerprint_call(_dummy_fn, (Twin(1, "a"),), {}, []) # Matches twin ID + + # 2. Call register_memo_key_function shortcut: delegates to key_fn AND preserves stable_id + register_memo_key_function(Entry, lambda e: ("key", e.value)) + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) == fingerprint_call( + _dummy_fn, (Entry(1, "b"),), {}, [] + ) # key_fn active (ignores 'ignored') + + # 3. Call unregister_memo_key_function shortcut: removes key_fn AND preserves stable_id + unregister_memo_key_function(Entry) + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) != fingerprint_call( + _dummy_fn, (Entry(1, "b"),), {}, [] + ) # key_fn removed: dataclass hashing includes 'ignored' + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) == fingerprint_call( + _dummy_fn, (Twin(1, "a"),), {}, [] + ) # stable_id still preserved and matches twin + + # 4. Full unregister clears everything: + unregister_memo_type(Entry) + assert fingerprint_call( + _dummy_fn, (Entry(1, "a"),), {}, [] + ) != fingerprint_call(_dummy_fn, (Twin(1, "a"),), {}, []) # stable_id removed + finally: + unregister_memo_type(Entry) + unregister_memo_type(Twin) + + +def test_unregister_memo_type_clears_key_function_and_stable_type_id() -> None: + class RegisteredOnly: + def __init__(self, value: object) -> None: + self.value = value + + class SameStableTypeId: + def __init__(self, value: object) -> None: + self.value = value + + try: + register_memo_type( + RegisteredOnly, + lambda entry: ("registered", entry.value), + stable_type_id="test.UnregisterCombined/v1", + ) + register_memo_type( + SameStableTypeId, + lambda entry: ("registered", entry.value), + stable_type_id="test.UnregisterCombined/v1", + ) + assert fingerprint_call( + _dummy_fn, (RegisteredOnly(1),), {}, [] + ) == fingerprint_call(_dummy_fn, (SameStableTypeId(1),), {}, []) + assert fingerprint_call(_dummy_fn, (RegisteredOnly,), {}, []) == ( + fingerprint_call(_dummy_fn, (SameStableTypeId,), {}, []) + ) + + unregister_memo_type(RegisteredOnly) + with pytest.raises(TypeError, match="Unsupported type for memoization key"): + fingerprint_call(_dummy_fn, (RegisteredOnly(1),), {}, []) + assert fingerprint_call(_dummy_fn, (RegisteredOnly,), {}, []) != ( + fingerprint_call(_dummy_fn, (SameStableTypeId,), {}, []) + ) + finally: + unregister_memo_type(RegisteredOnly) + unregister_memo_type(SameStableTypeId) + + class OldEntry: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class NewEntry: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + try: + register_memo_type(OldEntry, stable_type_id="test.UnregisterStable/v1") + register_memo_type(NewEntry, stable_type_id="test.UnregisterStable/v1") + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry(1),), {}, [] + ) + unregister_memo_type(OldEntry) + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) != fingerprint_call( + _dummy_fn, (NewEntry(1),), {}, [] + ) + finally: + unregister_memo_type(OldEntry) + unregister_memo_type(NewEntry) + + +def test_none_declared_stable_type_id_falls_back_to_registered_or_module_qualname() -> ( + None +): + class NoneId: + __coco_memo_type_id__ = None + + def __coco_memo_key__(self) -> object: + return ("bad", 1) + + none_id_canonical = _memo_fingerprint._canonicalize(NoneId(), None, []) + assert isinstance(none_id_canonical, tuple) + assert none_id_canonical[:3] == ( + "hook", + _memo_fingerprint.canonical_module_name(NoneId), + NoneId.__qualname__, + ) + + stable_type_id = "test.NoneDeclarationFallback/v1" + try: + register_memo_type(NoneId, stable_type_id=stable_type_id) + registered_canonical = _memo_fingerprint._canonicalize(NoneId(), None, []) + assert isinstance(registered_canonical, tuple) + assert registered_canonical[:3] == ( + "hook", + ("__coco_memo_type_id__", stable_type_id), + None, + ) + finally: + unregister_memo_type(NoneId) + + +def test_non_string_declared_stable_type_id_falls_back_to_registered_or_module_qualname() -> ( + None +): + class NonStringSourceEntry: + __coco_memo_type_id__ = 12345 + + def __coco_memo_key__(self) -> object: + return ("bad", 1) + + non_str_canonical = _memo_fingerprint._canonicalize( + NonStringSourceEntry(), None, [] + ) + assert isinstance(non_str_canonical, tuple) + assert non_str_canonical[:3] == ( + "hook", + _memo_fingerprint.canonical_module_name(NonStringSourceEntry), + NonStringSourceEntry.__qualname__, + ) + + stable_type_id = "test.NonStringDeclarationFallback/v1" + try: + register_memo_type(NonStringSourceEntry, stable_type_id=stable_type_id) + registered_canonical = _memo_fingerprint._canonicalize( + NonStringSourceEntry(), None, [] + ) + assert isinstance(registered_canonical, tuple) + assert registered_canonical[:3] == ( + "hook", + ("__coco_memo_type_id__", stable_type_id), + None, + ) + finally: + unregister_memo_type(NonStringSourceEntry) + + +def test_register_memo_type_validation_and_public_export() -> None: + import cocoindex as coco + + assert coco.register_memo_type is _memo_fingerprint.register_memo_type + assert ( + coco.register_memo_key_function is _memo_fingerprint.register_memo_key_function + ) + assert coco.prev_type_id is _memo_fingerprint.prev_type_id + assert "register_memo_type" in coco.__all__ + assert "register_memo_key_function" in coco.__all__ + assert "prev_type_id" in coco.__all__ + previous_type_id = coco.prev_type_id("old_package.models", "SourceEntry") + assert isinstance(previous_type_id, str) + assert not hasattr(coco, "register_memo_type_identifier") + assert not hasattr(coco, "register_not_memo_keyable") + assert not hasattr(coco, "unregister_memo_type") + assert not hasattr(coco, "unregister_memo_key_function") + with pytest.raises(TypeError, match="expects typ to be a type"): + coco.register_memo_type(cast(Any, object()), stable_type_id="test.Invalid/v1") + + class ExplicitNoneStateEntry: + pass + + try: + coco.register_memo_type( + ExplicitNoneStateEntry, + state_fn=None, + stable_type_id="test.ExplicitNoneState/v1", + ) + assert _memo_fingerprint._type_identity_parts( + ExplicitNoneStateEntry, + _memo_fingerprint._registered_memo_type_registry(ExplicitNoneStateEntry), + ) == (("__coco_memo_type_id__", "test.ExplicitNoneState/v1"), None) + finally: + _memo_fingerprint.unregister_memo_type(ExplicitNoneStateEntry) + + def test_cycles_are_supported_and_deterministic() -> None: # Self-cycle list a: Any = [] @@ -543,8 +2099,8 @@ def __init__(self, v: object) -> None: def _state_fn(obj: Any, prev: Any) -> MemoStateOutcome: return MemoStateOutcome(state=prev, memo_valid=True) - register_memo_key_function(Registered, lambda r: r.v, state_fn=_state_fn) try: + register_memo_key_function(Registered, lambda r: r.v, state_fn=_state_fn) methods: list[Any] = [] fingerprint_call(_dummy_fn, (Registered(7),), {}, state_methods=methods) assert len(methods) == 1