From 8ff244f1d7a7f5c008b84c9387239696fc3fd368 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 4 Jul 2026 12:48:01 -0700 Subject: [PATCH 01/16] feat(memo): use stable custom IDs for memo type identity --- python/cocoindex/_internal/api.py | 2 + .../cocoindex/_internal/memo_fingerprint.py | 232 ++++++++++++++++-- 2 files changed, 215 insertions(+), 19 deletions(-) diff --git a/python/cocoindex/_internal/api.py b/python/cocoindex/_internal/api.py index ca97a3a1a..21ba6e6d5 100644 --- a/python/cocoindex/_internal/api.py +++ b/python/cocoindex/_internal/api.py @@ -104,6 +104,7 @@ from .memo_fingerprint import ( memo_fingerprint, register_memo_key_function, + register_memo_type_identifier, NotMemoKeyable, ) @@ -944,6 +945,7 @@ class Cursor: # .memo_fingerprint "memo_fingerprint", "register_memo_key_function", + "register_memo_type_identifier", "NotMemoKeyable", # .pending_marker "MaybePendingS", diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index 07c9c65b6..34ee26896 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -10,12 +10,14 @@ import dataclasses import functools +import inspect import math import os import pickle import struct import sys import typing +import weakref from . import core from .serde import ( @@ -28,6 +30,7 @@ _KeyFn = typing.Callable[[typing.Any], typing.Any] +_BoundKeyFn = typing.Callable[[], typing.Any] _StateFn = typing.Callable[[typing.Any, typing.Any], typing.Any] @@ -37,6 +40,7 @@ class _MemoFns(typing.NamedTuple): _memo_fns: dict[type, _MemoFns] = {} +_memo_type_identifiers: dict[int, tuple[weakref.ReferenceType[type], str]] = {} class StateFnEntry(typing.NamedTuple): @@ -118,6 +122,192 @@ def canonical_module_name(obj: typing.Any) -> str: return mod +def _memo_type_label(typ: type) -> str: + """Return a user-facing label for type-identifier validation errors.""" + return f"{canonical_module_name(typ)}.{getattr(typ, '__qualname__', '')}" + + +def _validate_memo_type_identifier(identifier: object, *, source: str) -> str: + """Validate a non-empty stable memo type identifier.""" + if not isinstance(identifier, str): + raise TypeError(f"{source} must be a str, got {type(identifier).__name__}") + if identifier.strip() == "": + raise ValueError( + f"{source} must be non-empty and contain non-whitespace characters" + ) + return identifier + + +def _remove_memo_type_identifier( + type_id: int, dead_ref: weakref.ReferenceType[type] +) -> None: + """Remove ``type_id`` only if ``dead_ref`` is still the stored weakref.""" + entry = _memo_type_identifiers.get(type_id) + if entry is not None and entry[0] is dead_ref: + _memo_type_identifiers.pop(type_id, None) + + +def register_memo_type_identifier(typ: type, identifier: str) -> None: + """Register a stable memo type identity for one exact Python type. + + Type-aware fingerprints use it instead of module+qualname. Registration + overrides ``typ.__coco_memo_type_id__``; explicit class-object hooks and + pickle fallback must carry their own stable key. + """ + if not isinstance(typ, type): + raise TypeError( + "register_memo_type_identifier() expects typ to be a type, " + f"got {type(typ).__name__}" + ) + identifier = _validate_memo_type_identifier( + identifier, source="register_memo_type_identifier(..., identifier)" + ) + type_id = id(typ) + + def _remove_stale_type_identifier( + dead_ref: weakref.ReferenceType[type], + ) -> None: + _remove_memo_type_identifier(type_id, dead_ref) + + _memo_type_identifiers[type_id] = ( + weakref.ref(typ, _remove_stale_type_identifier), + identifier, + ) + + +def _unregister_memo_type_identifier(typ: type) -> None: + """Best-effort test helper for removing an exact-type registration.""" + type_id = id(typ) + entry = _memo_type_identifiers.get(type_id) + if entry is not None and entry[0]() is typ: + _memo_type_identifiers.pop(type_id, None) + + +def _registered_memo_type_identifier(typ: type) -> str | None: + """Return the registered identifier for ``typ`` from the id-keyed table.""" + type_id = id(typ) + entry = _memo_type_identifiers.get(type_id) + if entry is None: + return None + ref, identifier = entry + if ref() is typ: + return identifier + _memo_type_identifiers.pop(type_id, None) + return None + + +def _lookup_memo_type_identifier(typ: type) -> str | None: + """Resolve a registered or exact ``__coco_memo_type_id__`` identifier.""" + identifier = _registered_memo_type_identifier(typ) + if identifier is not None: + return identifier + if "__coco_memo_type_id__" in typ.__dict__: + return _validate_memo_type_identifier( + typ.__dict__["__coco_memo_type_id__"], + source=f"{_memo_type_label(typ)}.__coco_memo_type_id__", + ) + return None + + +_MEMO_KEY_ATTR = "__coco_memo_key__" +_MEMO_STATE_ATTR = "__coco_memo_state__" + + +def _callable_memo_hook(hook: object) -> _BoundKeyFn | None: + """Return ``hook`` when it is callable.""" + if not callable(hook): + return None + return typing.cast(_BoundKeyFn, hook) + + +def _python_callable_accepts_no_args(hook: object) -> bool: + """Return whether a Python function/method can bind zero arguments. + + This is deliberately limited to regular Python callables. Some valid + callable objects and C-extension functions cannot expose an inspectable + signature; those are accepted and validated by the actual call. + """ + if not (inspect.isfunction(hook) or inspect.ismethod(hook)): + return True + try: + inspect.signature(hook).bind() + except TypeError: + return False + return True + + +def _metaclass_object_memo_hook(cls: type) -> _BoundKeyFn | None: + """Resolve a memo key hook from ``cls``'s metaclass.""" + for base in typing.cast(type, type(cls)).__mro__: + raw = base.__dict__.get(_MEMO_KEY_ATTR) + if raw is None: + continue + hook = ( + typing.cast(typing.Any, raw).__get__(cls, type(cls)) + if hasattr(raw, "__get__") + else raw + ) + return _callable_memo_hook(hook) + return None + + +def _class_object_memo_hook(cls: type) -> _BoundKeyFn | None: + """Resolve an explicit memo key hook for a class object. + + Plain class-body methods that require an instance are ignored because there + is no ``self`` to bind. Zero-argument functions assigned as class + attributes, descriptor-based hooks, callable objects, and metaclass hooks + remain valid class-object hooks. + """ + for base in cls.__mro__: + raw = base.__dict__.get(_MEMO_KEY_ATTR) + if raw is None: + continue + hook = getattr(cls, _MEMO_KEY_ATTR, None) + if isinstance(raw, (classmethod, staticmethod)): + if not callable(hook) or not _python_callable_accepts_no_args(hook): + hook_kind = ( + "classmethod" if isinstance(raw, classmethod) else "staticmethod" + ) + raise TypeError( + f"{_memo_type_label(base)}.{_MEMO_KEY_ATTR} is a {hook_kind} " + "that cannot be called with zero arguments; class-object hooks " + "must take no arguments after binding" + ) + return typing.cast(_BoundKeyFn, hook) + if callable(hook) and _python_callable_accepts_no_args(hook): + return typing.cast(_BoundKeyFn, hook) + break + return _metaclass_object_memo_hook(cls) + + +def _class_object_state_fn_entry(cls: type) -> StateFnEntry | None: + """Resolve a memo state hook bound to ``cls``'s metaclass.""" + typ = type(cls) + for base in typing.cast(type, typ).__mro__: + raw = base.__dict__.get(_MEMO_STATE_ATTR) + if raw is None: + continue + state_hook = ( + typing.cast(typing.Any, raw).__get__(cls, typ) + if hasattr(raw, "__get__") + else raw + ) + if not callable(state_hook): + return None + raw_fn = getattr(typ, _MEMO_STATE_ATTR) + return _make_state_fn_entry(state_hook, raw_fn) + return None + + +def _type_identity_parts(typ: type) -> tuple[Fingerprintable, Fingerprintable]: + """Return stable-ID or module+qualname type identity parts.""" + identifier = _lookup_memo_type_identifier(typ) + if identifier is not None: + return (("__coco_memo_type_id__", identifier), None) + return (canonical_module_name(typ), getattr(typ, "__qualname__", None)) + + 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) @@ -142,8 +332,7 @@ def _canonicalize_dataclass( fields = dataclasses.fields(obj) # type: ignore[arg-type] return ( "dataclass", - canonical_module_name(typ), - typ.__qualname__, + *_type_identity_parts(typ), tuple( (field.name, _canonicalize(getattr(obj, field.name), _seen, state_methods)) for field in fields @@ -165,8 +354,7 @@ 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), tuple( (name, _canonicalize(getattr(obj, name), _seen, state_methods)) for name in field_names @@ -284,21 +472,30 @@ def _canonicalize( return bytes(obj) # 2) Hook / registry (apply once, then recurse on returned key fragment) - hook = getattr(obj, "__coco_memo_key__", None) + hook = ( + _class_object_memo_hook(obj) + if isinstance(obj, type) + else getattr(obj, _MEMO_KEY_ATTR, None) + ) if hook is not None and callable(hook): k = hook() typ = type(obj) tag = "hook" - state_hook = getattr(obj, "__coco_memo_state__", 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__") - state_methods.append(_make_state_fn_entry(state_hook, raw_fn)) + if isinstance(obj, type): + state_entry = _class_object_state_fn_entry(obj) + if state_entry is not None: + tag = "shook" + state_methods.append(state_entry) + else: + 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, _MEMO_STATE_ATTR) + state_methods.append(_make_state_fn_entry(state_hook, raw_fn)) return ( tag, - canonical_module_name(typ), - typ.__qualname__, + *_type_identity_parts(typ), _canonicalize(k, _seen, state_methods), ) @@ -313,8 +510,7 @@ def _canonicalize( state_methods.append(_make_state_fn_entry(bound, memo.state_fn)) return ( tag, - canonical_module_name(base), - base.__qualname__, + *_type_identity_parts(base), _canonicalize(k, _seen, state_methods), ) @@ -441,15 +637,13 @@ def fingerprint_call( # Register memo key for class types. -register_memo_key_function( - type, - lambda cls: (canonical_module_name(cls), getattr(cls, "__qualname__", None)), -) +register_memo_key_function(type, lambda cls: _type_identity_parts(cls)) __all__ = [ "NotMemoKeyable", "register_memo_key_function", + "register_memo_type_identifier", "register_not_memo_keyable", "unregister_memo_key_function", "fingerprint_call", From 29e4de0d0b63892b932b549014a7d33ecb0251ca Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 4 Jul 2026 12:48:04 -0700 Subject: [PATCH 02/16] test(memo): prove stable IDs survive renames and module moves --- python/tests/core/test_function_memo.py | 68 +++ .../tests/internal/test_memo_fingerprint.py | 431 +++++++++++++++++- 2 files changed, 498 insertions(+), 1 deletion(-) diff --git a/python/tests/core/test_function_memo.py b/python/tests/core/test_function_memo.py index 4a5b73b8d..e56541929 100644 --- a/python/tests/core/test_function_memo.py +++ b/python/tests/core/test_function_memo.py @@ -27,6 +27,24 @@ 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 class DictSourceDataEntry: name: str @@ -38,6 +56,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 +75,54 @@ 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_memo_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_memo_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_memo_pure_function() -> None: GlobalDictTarget.store.clear() _plain_source_data.clear() diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index ab7ca2f32..e8a7a1670 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -1,12 +1,15 @@ import dataclasses import math -from typing import Any +from typing import Any, ClassVar, cast import pytest from cocoindex._internal.function import _apply_memo_key, _normalize_memo_key from cocoindex._internal.memo_fingerprint import ( fingerprint_call, + StateFnEntry, + _unregister_memo_type_identifier, + register_memo_type_identifier, register_memo_key_function, unregister_memo_key_function, ) @@ -132,6 +135,432 @@ 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_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),), {}, []) + + +class _UninspectableZeroArgMemoHook: + @property + def __signature__(self) -> object: + raise ValueError("signature unavailable") + + def __call__(self) -> object: + return ("class-hook", "stable") + + +def test_raw_class_object_honors_zero_arg_memo_key() -> None: + def stable_key() -> object: + return ("class-hook", "stable") + + def other_key() -> object: + return ("class-hook", "other") + + class OldEntry: + __coco_memo_key__ = stable_key + + class NewEntry: + __coco_memo_key__ = stable_key + + class ChangedEntry: + __coco_memo_key__ = other_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, (NewEntry,), {}, []) != fingerprint_call( + _dummy_fn, (ChangedEntry,), {}, [] + ) + + +def test_raw_class_object_metaclass_hook_collects_state_method() -> None: + class MemoMeta(type): + def __coco_memo_key__(cls) -> object: + return ("metaclass-hook", "stable") + + def __coco_memo_state__(cls, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome( + state=("state", cls.__name__, prev_state), memo_valid=True + ) + + class OldEntry(metaclass=MemoMeta): + pass + + class NewEntry(metaclass=MemoMeta): + pass + + OldEntry.__module__ = "tests.old_raw_class" + NewEntry.__module__ = "tests.new_raw_class" + + old_methods: list[Any] = [] + new_methods: list[Any] = [] + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, old_methods) == fingerprint_call( + _dummy_fn, (NewEntry,), {}, new_methods + ) + + assert len(old_methods) == 1 + assert len(new_methods) == 1 + assert isinstance(old_methods[0], StateFnEntry) + assert isinstance(new_methods[0], StateFnEntry) + assert old_methods[0].call("previous") == MemoStateOutcome( + state=("state", "OldEntry", "previous"), memo_valid=True + ) + assert new_methods[0].call("previous") == MemoStateOutcome( + state=("state", "NewEntry", "previous"), memo_valid=True + ) + + +def test_raw_class_object_accepts_uninspectable_zero_arg_staticmethod_memo_key() -> None: + hook = _UninspectableZeroArgMemoHook() + + class OldEntry: + __coco_memo_key__ = staticmethod(hook) + + class NewEntry: + __coco_memo_key__ = staticmethod(hook) + + OldEntry.__module__ = "tests.old_raw_class" + NewEntry.__module__ = "tests.new_raw_class" + + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry,), {}, [] + ) + + +def test_raw_class_object_accepts_uninspectable_zero_arg_plain_memo_key() -> None: + hook = _UninspectableZeroArgMemoHook() + + class OldEntry: + __coco_memo_key__ = hook + + class NewEntry: + __coco_memo_key__ = hook + + OldEntry.__module__ = "tests.old_raw_class" + NewEntry.__module__ = "tests.new_raw_class" + + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry,), {}, [] + ) + + +def test_raw_class_object_rejects_non_zero_arg_staticmethod_memo_key() -> None: + """Reject invalid staticmethod hooks before calling class-object memo keys.""" + class Entry: + @staticmethod + def __coco_memo_key__(value: object) -> object: + return ("class-hook", value) + + with pytest.raises( + TypeError, + match=( + r"Entry\.__coco_memo_key__ is a staticmethod that cannot be called " + r"with zero arguments; class-object hooks must take no arguments " + r"after binding" + ), + ): + fingerprint_call(_dummy_fn, (Entry,), {}, []) + + +def test_raw_class_object_rejects_non_zero_arg_classmethod_memo_key() -> None: + """Reject invalid classmethod hooks after descriptor binding leaves arguments.""" + class Entry: + @classmethod + def __coco_memo_key__(cls, value: object) -> object: + return ("class-hook", cls.__name__, value) + + with pytest.raises( + TypeError, + match=( + r"Entry\.__coco_memo_key__ is a classmethod that cannot be called " + r"with zero arguments; class-object hooks must take no arguments " + r"after binding" + ), + ): + fingerprint_call(_dummy_fn, (Entry,), {}, []) + + +def test_raw_class_object_ignores_instance_memo_state() -> None: + class Entry: + @staticmethod + def __coco_memo_key__() -> object: + return ("class-hook", "stable") + + def __coco_memo_state__(self, prev_state: object) -> MemoStateOutcome: + raise AssertionError("instance memo state must not run for class objects") + + methods: list[Any] = [] + fp1 = fingerprint_call(_dummy_fn, (Entry,), {}, methods) + fp2 = fingerprint_call(_dummy_fn, (Entry,), {}, []) + + assert fp1 == fp2 + assert methods == [] + + +def test_raw_class_object_uses_metaclass_hook_after_ignored_instance_key() -> None: + class MemoMeta(type): + def __coco_memo_key__(cls) -> object: + return ("metaclass-hook", "stable") + + class Base: + def __coco_memo_key__(self) -> object: + raise AssertionError("instance memo key must not run for class objects") + + class OldEntry(Base, metaclass=MemoMeta): + pass + + class NewEntry(Base, metaclass=MemoMeta): + pass + + OldEntry.__module__ = "tests.old_raw_class" + NewEntry.__module__ = "tests.new_raw_class" + + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry,), {}, [] + ) + + +def test_raw_class_object_stable_type_id_ignores_instance_memo_key() -> None: + class OldEntry: + __coco_memo_type_id__ = "test.RawClass/v1" + + def __coco_memo_key__(self) -> object: + raise AssertionError("instance memo key must not run for class objects") + + class NewEntry: + __coco_memo_type_id__ = "test.RawClass/v1" + + def __coco_memo_key__(self) -> object: + raise AssertionError("instance memo key must not run for class objects") + + class ChangedEntry: + __coco_memo_type_id__ = "test.RawClass/v2" + + def __coco_memo_key__(self) -> object: + raise AssertionError("instance memo key must not run for class objects") + + 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_raw_class_object_explicit_hook_owns_stable_namespace() -> None: + """Explicit class-object hooks are custom keys, not stable-ID wrappers.""" + def shared_key() -> object: + return ("class-hook", "shared") + + def changed_key() -> object: + return ("class-hook", "changed") + + class OldEntry: + __coco_memo_type_id__ = "test.RawClassHook/v1" + __coco_memo_key__ = staticmethod(shared_key) + + class ChangedKeyEntry: + __coco_memo_type_id__ = "test.RawClassHook/v1" + __coco_memo_key__ = staticmethod(changed_key) + + class ChangedTypeIdEntry: + __coco_memo_type_id__ = "test.RawClassHook/v2" + __coco_memo_key__ = staticmethod(shared_key) + + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) != fingerprint_call( + _dummy_fn, (ChangedKeyEntry,), {}, [] + ) + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( + _dummy_fn, (ChangedTypeIdEntry,), {}, [] + ) + + +def test_registered_memo_type_identifier_allows_renamed_type_reuse() -> 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) + + register_memo_type_identifier(OldEntry, "test.RegisteredEntry/v1") + register_memo_type_identifier(NewEntry, "test.RegisteredEntry/v1") + try: + 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_identifier(OldEntry) + _unregister_memo_type_identifier(NewEntry) + + +def test_registered_memo_key_function_uses_registered_base_type_identity() -> None: + class Base: + def __init__(self, value: object) -> None: + self.value = value + + class ChildA(Base): + pass + + class ChildB(Base): + pass + + register_memo_key_function(Base, lambda entry: ("base", entry.value)) + register_memo_type_identifier(Base, "test.RegisteredBase/v1") + try: + 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_key_function(Base) + _unregister_memo_type_identifier(Base) + + +def test_registered_memo_type_identifier_is_identity_exact_for_equal_metaclasses() -> ( + None +): + class EqMeta(type): + def __eq__(cls, other: object) -> bool: + return isinstance(other, EqMeta) + + def __hash__(cls) -> int: + return 1 + + class A(metaclass=EqMeta): + def __coco_memo_key__(self) -> object: + return ("same",) + + class B(metaclass=EqMeta): + def __coco_memo_key__(self) -> object: + return ("same",) + + assert fingerprint_call(_dummy_fn, (A(),), {}, []) != fingerprint_call( + _dummy_fn, (B(),), {}, [] + ) + + register_memo_type_identifier(A, "test.EqualityMetaA/v1") + try: + assert fingerprint_call(_dummy_fn, (A(),), {}, []) != fingerprint_call( + _dummy_fn, (B(),), {}, [] + ) + finally: + _unregister_memo_type_identifier(A) + + +def test_stable_type_id_exact_type_and_validation() -> None: + class Parent: + __coco_memo_type_id__ = "test.Parent/v1" + + def __coco_memo_key__(self) -> object: + return ("same", 1) + + class Child(Parent): + pass + + class BadObjectId: + __coco_memo_type_id__ = object() + + def __coco_memo_key__(self) -> object: + return ("bad", 1) + + class EmptyId: + __coco_memo_type_id__ = "" + + def __coco_memo_key__(self) -> object: + return ("bad", 1) + + assert fingerprint_call(_dummy_fn, (Parent(),), {}, []) != fingerprint_call( + _dummy_fn, (Child(),), {}, [] + ) + with pytest.raises(TypeError, match="must be a str"): + fingerprint_call(_dummy_fn, (BadObjectId(),), {}, []) + with pytest.raises(ValueError, match="non-empty"): + fingerprint_call(_dummy_fn, (EmptyId(),), {}, []) + + +def test_register_memo_type_identifier_validation_and_export() -> None: + import cocoindex as coco + + class Entry: + pass + + assert coco.register_memo_type_identifier is register_memo_type_identifier + with pytest.raises(TypeError, match="expects typ to be a type"): + register_memo_type_identifier(cast(Any, object()), "test.Invalid/v1") + with pytest.raises(TypeError, match="must be a str"): + register_memo_type_identifier(Entry, cast(Any, object())) + with pytest.raises(ValueError, match="non-empty"): + register_memo_type_identifier(Entry, "") + with pytest.raises(ValueError, match="non-empty"): + register_memo_type_identifier(Entry, " ") + + def test_cycles_are_supported_and_deterministic() -> None: # Self-cycle list a: Any = [] From 5335e477b1c813daef6c688f73119de76ef4613c Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 4 Jul 2026 12:48:06 -0700 Subject: [PATCH 03/16] docs(memo): describe stable type ID APIs and caveats --- .../docs/advanced_topics/memoization_keys.mdx | 44 +++++++++++++++++-- 1 file changed, 40 insertions(+), 4 deletions(-) diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index 0b59b8223..d94fe607c 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -1,14 +1,15 @@ --- 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 functions, or stable semantic type identifiers, 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: - **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 identifiers** — how to keep a type's memo namespace stable across refactors. ## How data fingerprinting works @@ -24,7 +25,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 semantic 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). @@ -78,6 +79,41 @@ register_memo_key_function(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). + +### Stabilize type identity across refactors + +By default, CocoIndex includes a type's module and qualified name around type-aware values such as dataclasses, Pydantic models, values with `__coco_memo_key__`, registered memo-key functions, and class objects. If a class moves or is renamed, that default identity changes and old memo entries are not reused. + +Add a stable semantic identifier when the type should keep the same memo namespace across refactors: + +```python +from dataclasses import dataclass +from typing import ClassVar + +@dataclass +class SourceEntry: + __coco_memo_type_id__: ClassVar[str] = "com.example.SourceEntry/v1" + + id: str + content: str +``` + +For a third-party type that is already handled by `register_memo_key_function(...)`, register the type ID on that same owner type: + +```python +from pathlib import Path +import cocoindex as coco + +coco.register_memo_key_function(Path, lambda path: str(path)) +coco.register_memo_type_identifier(Path, "stdlib.pathlib.Path/v1") +``` + +Direct stable type IDs are exact-type: `__coco_memo_type_id__` and direct `register_memo_type_identifier(SomeType, ...)` lookup apply only to that exact type. For values handled by `register_memo_key_function(...)`, the selected registered owner type supplies the type identity for subclasses, matching the MRO-aware key-function lookup described above. + +Stable type IDs do not retroactively migrate memo entries created before the ID existed, and they do not make plain pickle-fallback instances refactor-stable; use `__coco_memo_key__` or `register_memo_key_function(...)` when the object's value also needs a stable custom key. + +Class objects have one extra boundary: if no explicit class-object hook is selected, the registered `type` memo-key path can use the class object's stable type ID. An explicit zero-argument `__coco_memo_key__` hook on a class object is a custom key, so CocoIndex does not automatically wrap that hook payload in the class object's stable ID; include any stable namespace in the hook return value when needed. + ### 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). From 53d8defe95c06d1c46e79731822f93254931ab39 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 4 Jul 2026 22:32:30 -0700 Subject: [PATCH 04/16] test(memo): cover class object state hooks --- .../tests/core/test_memo_state_validation.py | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/python/tests/core/test_memo_state_validation.py b/python/tests/core/test_memo_state_validation.py index aa7df404a..2d697e00c 100644 --- a/python/tests/core/test_memo_state_validation.py +++ b/python/tests/core/test_memo_state_validation.py @@ -713,3 +713,111 @@ def test_state_changed_but_reusable_component_sync() -> None: ) app.update_blocking() assert _metrics.collect() == {"call.declare_two_level": 1} + + + +# ============================================================================ +# Raw class-object metaclass state validation (sync) +# ============================================================================ + +_class_object_source: dict[str, type] = {} +_class_object_prev_states: list[Any] = [] + + +class _ClassObjectMemoMeta(type): + def __coco_memo_key__(cls) -> object: + return getattr(cls, "stable_key") + + def __coco_memo_state__(cls, prev_state: Any) -> coco.MemoStateOutcome: + state_value = getattr(cls, "state_value") + _class_object_prev_states.append(prev_state) + memo_valid = not coco.is_non_existence(prev_state) and prev_state == state_value + return coco.MemoStateOutcome(state=state_value, memo_valid=memo_valid) + + +def _make_class_object( + stable_key: str, state_value: int, content: str +) -> type: + return _ClassObjectMemoMeta( + f"ClassObject_{stable_key}_{state_value}_{content}", + (), + { + "stable_key": stable_key, + "state_value": state_value, + "content": content, + }, + ) + + +@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_state_validation_sync() -> None: + """Class-object metaclass state participates in App memo validation.""" + GlobalDictTarget.store.clear() + _class_object_source.clear() + _class_object_prev_states.clear() + _metrics.clear() + + app = coco.App( + coco.AppConfig( + name="test_class_object_metaclass_state_validation_sync", + environment=coco_env, + ), + _process_class_objects, + ) + + # Run 1: no previous state, so the transform executes and persists content A. + _class_object_source["row"] = _make_class_object("stable", 1, "A") + app.update_blocking() + assert _metrics.collect() == {"call.transform_class_object": 1} + assert len(_class_object_prev_states) == 1 + assert coco.is_non_existence(_class_object_prev_states[0]) + assert GlobalDictTarget.store.data == { + "row": DictDataWithPrev( + data="class content: A", prev=[], prev_may_be_missing=True + ) + } + + # Run 2: equivalent replacement with same state is a cache hit; C is ignored. + _class_object_source["row"] = _make_class_object("stable", 1, "C") + app.update_blocking() + assert _metrics.collect() == {} + assert len(_class_object_prev_states) == 2 + assert coco.is_non_existence(_class_object_prev_states[0]) + assert _class_object_prev_states[1:] == [1] + assert GlobalDictTarget.store.data["row"].data == "class content: A" + + # Run 3: same key but new state invalidates the cache and persists content B. + _class_object_source["row"] = _make_class_object("stable", 2, "B") + app.update_blocking() + assert _metrics.collect() == {"call.transform_class_object": 1} + assert len(_class_object_prev_states) == 3 + assert coco.is_non_existence(_class_object_prev_states[0]) + assert _class_object_prev_states[1:] == [1, 1] + assert GlobalDictTarget.store.data == { + "row": DictDataWithPrev( + data="class content: B", + prev=["class content: A"], + prev_may_be_missing=False, + ) + } + + # Run 4: state 2 was persisted, so another replacement hits cache; D is ignored. + _class_object_source["row"] = _make_class_object("stable", 2, "D") + app.update_blocking() + assert _metrics.collect() == {} + assert len(_class_object_prev_states) == 4 + assert coco.is_non_existence(_class_object_prev_states[0]) + assert _class_object_prev_states[1:] == [1, 1, 2] + assert GlobalDictTarget.store.data["row"].data == "class content: B" \ No newline at end of file From 985212a3795dca1a823409350abe1eb541f6882b Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sun, 5 Jul 2026 16:10:01 -0700 Subject: [PATCH 05/16] test: fix memo fingerprint prek formatting --- python/tests/core/test_memo_state_validation.py | 7 ++----- python/tests/internal/test_memo_fingerprint.py | 13 +++++++++---- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/python/tests/core/test_memo_state_validation.py b/python/tests/core/test_memo_state_validation.py index 2d697e00c..7e47c6be4 100644 --- a/python/tests/core/test_memo_state_validation.py +++ b/python/tests/core/test_memo_state_validation.py @@ -715,7 +715,6 @@ def test_state_changed_but_reusable_component_sync() -> None: assert _metrics.collect() == {"call.declare_two_level": 1} - # ============================================================================ # Raw class-object metaclass state validation (sync) # ============================================================================ @@ -735,9 +734,7 @@ def __coco_memo_state__(cls, prev_state: Any) -> coco.MemoStateOutcome: return coco.MemoStateOutcome(state=state_value, memo_valid=memo_valid) -def _make_class_object( - stable_key: str, state_value: int, content: str -) -> type: +def _make_class_object(stable_key: str, state_value: int, content: str) -> type: return _ClassObjectMemoMeta( f"ClassObject_{stable_key}_{state_value}_{content}", (), @@ -820,4 +817,4 @@ def test_class_object_metaclass_state_validation_sync() -> None: assert len(_class_object_prev_states) == 4 assert coco.is_non_existence(_class_object_prev_states[0]) assert _class_object_prev_states[1:] == [1, 1, 2] - assert GlobalDictTarget.store.data["row"].data == "class content: B" \ No newline at end of file + assert GlobalDictTarget.store.data["row"].data == "class content: B" diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index e8a7a1670..0cb264236 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -243,9 +243,9 @@ class NewEntry(metaclass=MemoMeta): old_methods: list[Any] = [] new_methods: list[Any] = [] - assert fingerprint_call(_dummy_fn, (OldEntry,), {}, old_methods) == fingerprint_call( - _dummy_fn, (NewEntry,), {}, new_methods - ) + assert fingerprint_call( + _dummy_fn, (OldEntry,), {}, old_methods + ) == fingerprint_call(_dummy_fn, (NewEntry,), {}, new_methods) assert len(old_methods) == 1 assert len(new_methods) == 1 @@ -259,7 +259,9 @@ class NewEntry(metaclass=MemoMeta): ) -def test_raw_class_object_accepts_uninspectable_zero_arg_staticmethod_memo_key() -> None: +def test_raw_class_object_accepts_uninspectable_zero_arg_staticmethod_memo_key() -> ( + None +): hook = _UninspectableZeroArgMemoHook() class OldEntry: @@ -295,6 +297,7 @@ class NewEntry: def test_raw_class_object_rejects_non_zero_arg_staticmethod_memo_key() -> None: """Reject invalid staticmethod hooks before calling class-object memo keys.""" + class Entry: @staticmethod def __coco_memo_key__(value: object) -> object: @@ -313,6 +316,7 @@ def __coco_memo_key__(value: object) -> object: def test_raw_class_object_rejects_non_zero_arg_classmethod_memo_key() -> None: """Reject invalid classmethod hooks after descriptor binding leaves arguments.""" + class Entry: @classmethod def __coco_memo_key__(cls, value: object) -> object: @@ -402,6 +406,7 @@ def __coco_memo_key__(self) -> object: def test_raw_class_object_explicit_hook_owns_stable_namespace() -> None: """Explicit class-object hooks are custom keys, not stable-ID wrappers.""" + def shared_key() -> object: return ("class-hook", "shared") From f6e694a75506fcdb9e6f04e6fe1df85eafa43e07 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Thu, 9 Jul 2026 14:19:24 -0700 Subject: [PATCH 06/16] feat(memo): implement reviewer feedback, fix bugs --- .../docs/advanced_topics/memoization_keys.mdx | 280 ++++- .../docs/programming_guide/function.mdx | 24 +- python/cocoindex/_internal/api.py | 2 - python/cocoindex/_internal/context_keys.py | 2 +- python/cocoindex/_internal/function.py | 18 +- .../cocoindex/_internal/memo_fingerprint.py | 515 ++++---- python/tests/core/test_function_memo.py | 4 +- .../tests/core/test_memo_state_validation.py | 129 +- .../tests/internal/test_memo_fingerprint.py | 1051 +++++++++++++---- 9 files changed, 1504 insertions(+), 521 deletions(-) diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index d94fe607c..850f992ab 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -2,22 +2,36 @@ title: "*Memoization keys* & states" description: > Customize how CocoIndex fingerprints memoized inputs via __coco_memo_key__, - registered functions, or stable semantic type identifiers, and layer state - validation (e.g. mtime, then content hash) on top. + 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: - **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 identifiers** — how to keep a type's memo namespace stable across refactors. +- **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: +Timing depends on where the value comes from. Function arguments are fingerprinted when the function is called. Change-detected context values are fingerprinted when they are provided (for example, by `builder.provide()`), before `use_context()` reads them. `deps` values are fingerprinted once when the `@coco.fn(..., deps=...)` or `@coco.fn.as_async(..., deps=...)` decorator is applied. -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. +For each data value, CocoIndex derives a canonical form with this precedence: + +1. **Class objects** (`type`) — the class values themselves, such as `ProductRow`, not instances such as `ProductRow(...)`. CocoIndex first checks for an explicit memo key registration on the class object's metaclass or its bases, including `type`. If none exists, it uses the built-in class-object identity path. It never calls memo-key or memo-state attributes on the class object itself. See the class-object notes below for what is and isn't included. +2. For non-class objects, if the object implements **`__coco_memo_key__()`**, CocoIndex uses its return value. +3. Otherwise, if you registered a **memo key function** for the object's type, CocoIndex uses that. +4. Otherwise, CocoIndex uses built-in canonicalization for supported automatic types, including primitives, containers, dataclasses, Pydantic models, and pickle fallback. The following types are handled automatically (no custom key needed): @@ -25,7 +39,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 semantic type ID when configured, otherwise 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). @@ -60,9 +74,22 @@ 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 @@ -79,12 +106,13 @@ register_memo_key_function(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 -### Stabilize type identity across refactors +This is an advanced migration feature. If you are not preserving existing memo entries across a class move or rename, skip this section and continue with [`memo_key=`](#override-at-the-call-site-with-memo_key) or [memo state validation](#memo-state-validation). -By default, CocoIndex includes a type's module and qualified name around type-aware values such as dataclasses, Pydantic models, values with `__coco_memo_key__`, registered memo-key functions, and class objects. If a class moves or is renamed, that default identity changes and old memo entries are not reused. +By default, CocoIndex includes a type namespace (module plus qualified name) for dataclass instances, Pydantic model instances, objects handled by `__coco_memo_key__` or a registered memo key function, and class objects passed as values. If a class moves or is renamed, that default type namespace changes and old memo entries are not reused. -Add a stable semantic identifier when the type should keep the same memo namespace across refactors: +A stable type ID replaces that type-name portion when the type should keep the same memo namespace across refactors: ```python from dataclasses import dataclass @@ -98,37 +126,184 @@ class SourceEntry: content: str ``` -For a third-party type that is already handled by `register_memo_key_function(...)`, register the type ID on that same owner type: +To preserve reuse **after** a refactor, add the stable type ID before moving or renaming the class and ensure every affected memoized call executes under that ID. This seeding execution is cold for affected values: the stable type ID replaces the old module-plus-qualified-name namespace, so CocoIndex creates new memo entries instead of reusing the old ones. + +An application update seeds only memoized calls that actually execute. A cached parent does not run its body or nested `@coco.fn` calls, so invalidate or reprocess affected parents—or otherwise invoke the affected nested calls—before the refactor. + +After seeding, move or rename the class while keeping the same stable type ID so later updates can reuse the entries created under that ID. Stable type IDs preserve an input type's identity across refactors; normal function logic tracking still invalidates memo entries when the decorated function's source changes. Adding the stable type ID only after the move does not reconstruct entries created under the old module plus qualified name. + +If you control the class definition, set `__coco_memo_type_id__` on the class. If you cannot or do not want to edit the class definition, register the same stable type ID with `register_memo_key_function`. + +Use `stable_type_id` without a key function only when CocoIndex already fingerprints the exact runtime type: dataclass instances, Pydantic model instances, instances with `__coco_memo_key__()`, or the class object itself. Subclasses do not inherit this form; give the subclass its own stable type ID unless a registered key function handles it through MRO lookup. + +Rule of thumb: if CocoIndex would otherwise use pickle for your object, `stable_type_id` by itself is not enough; register a key function and stable type ID together. + +Omit the key function argument entirely when registering only a stable type ID; do not pass `None` as the second argument. ```python -from pathlib import Path +from dataclasses import dataclass import cocoindex as coco -coco.register_memo_key_function(Path, lambda path: str(path)) -coco.register_memo_type_identifier(Path, "stdlib.pathlib.Path/v1") +@dataclass +class ProductRow: + sku: str + updated_at: int + +PRODUCT_ROW_TYPE_ID = "com.example.ProductRow/v1" + +coco.register_memo_key_function( + ProductRow, + stable_type_id=PRODUCT_ROW_TYPE_ID, +) ``` -Direct stable type IDs are exact-type: `__coco_memo_type_id__` and direct `register_memo_type_identifier(SomeType, ...)` lookup apply only to that exact type. For values handled by `register_memo_key_function(...)`, the selected registered owner type supplies the type identity for subclasses, matching the MRO-aware key-function lookup described above. +This registration is process-global. Register before the value is fingerprinted: for change-detected context values, before `builder.provide()`; for function arguments, before memoized calls that should use it; and for `deps`, before the `@coco.fn(..., deps=...)` or `@coco.fn.as_async(..., deps=...)` decorator executes. Later calls replace the previous registration for that exact type. -Stable type IDs do not retroactively migrate memo entries created before the ID existed, and they do not make plain pickle-fallback instances refactor-stable; use `__coco_memo_key__` or `register_memo_key_function(...)` when the object's value also needs a stable custom key. +This registration has two effects: -Class objects have one extra boundary: if no explicit class-object hook is selected, the registered `type` memo-key path can use the class object's stable type ID. An explicit zero-argument `__coco_memo_key__` hook on a class object is a custom key, so CocoIndex does not automatically wrap that hook payload in the class object's stable ID; include any stable namespace in the hook return value when needed. +- `ProductRow(...)` instances keep their normal dataclass field fingerprint, but the surrounding type namespace uses `PRODUCT_ROW_TYPE_ID`. +- The class object `ProductRow` also uses `PRODUCT_ROW_TYPE_ID` when you pass the class itself (`ProductRow`), not an instance (`ProductRow(...)`). -### Override at the call site with `memo_key=` +If the memoized function takes the class itself and reads class-level details, a stable type ID alone is not enough. Examples include field names, field types/defaults, validators, Pydantic config, class attributes, or version constants. Because `memo_key=` replaces the normal fingerprint for `row_type` in the example below, include both the stable type ID and the class-level details the function reads: + +```python +# Continuing the ProductRow example above. +from dataclasses import fields +import cocoindex as coco + +@coco.fn( + memo=True, + memo_key={ + "row_type": lambda row_type: ( + PRODUCT_ROW_TYPE_ID, + tuple(field.name for field in fields(row_type)), + ) + }, +) +def product_row_fields(row_type: type[ProductRow]) -> tuple[str, ...]: + return tuple(field.name for field in fields(row_type)) + +product_row_fields(ProductRow) +``` + +:::::caution[Class object schema is not fingerprinted automatically] +A stable type ID says the class should keep the same memo namespace after a move or rename. It does not say the class schema is unchanged. If cached results depend on class-level details, include those details with `memo_key=` for that function. If a class change should invalidate every memo that uses this stable type ID, change the stable type ID, for example from `/v1` to `/v2`. +::::: + +Use stable type IDs directly for class objects only when the memoized computation does not depend on class details such as module, qualified name, display name, schema, or state. Those details are intentionally excluded from the class object's fingerprint when a stable type ID is used. + +For a third-party type that would otherwise fall back to pickle, register both a custom key function and a stable type ID on the same owner type: -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). +```python +from some_library import ExternalRecord +import cocoindex as coco + +coco.register_memo_key_function( + ExternalRecord, + lambda record: (record.id, record.version), + stable_type_id="com.example.ExternalRecord/v1", +) +``` + +Registration is replacement-based, not additive. Each call replaces the full registration for that exact type: + +| Call shape | Registered key/state functions after the call | Registered stable type ID after the call | +| --- | --- | --- | +| `register_memo_key_function(T, key_fn, state_fn=..., stable_type_id="...")` — see [`state_fn` requirements](#register-a-state-function-when-you-dont-control-the-type) | `key_fn` and `state_fn` are installed. | The provided stable type ID is installed. | +| `register_memo_key_function(T, key_fn)` | `key_fn` is installed and any previous `state_fn` is cleared unless a new `state_fn` is supplied. | Any previous registered stable type ID is cleared. | +| `register_memo_key_function(T, stable_type_id="...")` | Any previous key/state functions are cleared. | The provided stable type ID is installed. | + +A `__coco_memo_type_id__` defined on the class is separate from the process-global registration table. If both are set for the same exact type, the registered stable type ID takes precedence. Provide `stable_type_id=` in a registration only when you intentionally want that override; otherwise, register only the key/state function and use the class-defined ID. + +For a registry-owned stable type ID, if a type needs a key function, state function, and stable type ID, register all of them in the same call. Avoid splitting registration across helper modules where import order can determine which registration wins. + +If a type already has a key or state function and you want to add a stable type ID, update the same call: ```python -@coco.fn(memo=True, memo_key={"entry": lambda e: (e.name, e.version), "extra": None}) -def transform(entry: SourceDataEntry, extra: str) -> str: +coco.register_memo_key_function( + T, + key_fn, + state_fn=state_fn, + stable_type_id="com.example.T/v1", +) +``` + +Do not add the stable type ID with a separate later call; that later call replaces the registration and clears the previous key/state functions. + +#### Stable type ID without a key function: exact type only + +`__coco_memo_type_id__` and `register_memo_key_function(SomeType, stable_type_id=...)` without a key function apply only to exactly `SomeType`, not subclasses. For example: + +```python +class Parent: + ... + +class Child(Parent): + ... + +coco.register_memo_key_function(Parent, stable_type_id="com.example.Parent/v1") +# Child(...) does not use Parent's stable type ID, because this registration +# has no key function and applies only to exactly Parent. +``` + +#### Stable type ID with a key function: owner selected by MRO + +When `register_memo_key_function(Parent, key_fn, stable_type_id=...)` handles a `Child(Parent)` value through MRO lookup, the selected registered owner (`Parent`) supplies the type namespace. That owner namespace can come from the registration's `stable_type_id=...` argument or from `Parent.__coco_memo_type_id__`. For example: + +```python +class Parent: + ... + +class Child(Parent): ... + +coco.register_memo_key_function( + Parent, + lambda value: value.id, + stable_type_id="com.example.Parent/v1", +) +# Child(...) can use Parent's stable type ID when Parent's key function wins +# MRO lookup and handles the Child value. +``` + +Stable type IDs do not retroactively migrate memo entries created before the stable type ID existed. + +#### Class objects as inputs + +Class objects are the class values themselves, such as `ProductRow`, not instances such as `ProductRow(...)`. + +When CocoIndex fingerprints a class object: + +- It never calls `__coco_memo_key__` or `__coco_memo_state__` attributes on the class object itself. +- It checks memo key functions explicitly registered on the class object's metaclass or its bases, including `type`. The most specific registered owner wins; a `state_fn` from the same registration participates in state validation, and that owner supplies the stable type ID or module-plus-qualified-name namespace. +- If no metaclass or `type` key function is registered, it uses the class object's own stable type ID when present, otherwise its module plus qualified name. +- A key/state registration on `ProductRow` customizes `ProductRow(...)` instances only. To customize the raw `ProductRow` class globally, register a key function on its metaclass. A stable-ID-only registration on a metaclass or on `type` does not customize all class objects. +- It excludes class schema and class state. If a memoized result depends on fields, field types/defaults, validators, Pydantic config, class attributes, or any other class-level details, pass that data separately with `memo_key=` or change the stable type ID. + +Do not add a `@classmethod` or `@staticmethod` named `__coco_memo_key__` to customize class-object memoization. Class objects ignore it, but instances can still see it and accidentally give every instance the same memo key. + +### Override at the call site with `memo_key=` + +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 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: @@ -139,11 +314,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 @@ -153,14 +323,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=...)`: @@ -195,7 +365,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): @@ -209,8 +381,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] @@ -225,30 +399,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_key_function`; `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_key_function( + 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 @@ -288,21 +474,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..7d3c8a4d8 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_key_function(..., 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 21ba6e6d5..ca97a3a1a 100644 --- a/python/cocoindex/_internal/api.py +++ b/python/cocoindex/_internal/api.py @@ -104,7 +104,6 @@ from .memo_fingerprint import ( memo_fingerprint, register_memo_key_function, - register_memo_type_identifier, NotMemoKeyable, ) @@ -945,7 +944,6 @@ class Cursor: # .memo_fingerprint "memo_fingerprint", "register_memo_key_function", - "register_memo_type_identifier", "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..35fc04114 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_key_function(..., 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 34ee26896..4c3c137bc 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -1,16 +1,15 @@ """ 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 import dataclasses import functools -import inspect import math import os import pickle @@ -30,17 +29,24 @@ _KeyFn = typing.Callable[[typing.Any], typing.Any] -_BoundKeyFn = typing.Callable[[], typing.Any] _StateFn = typing.Callable[[typing.Any, typing.Any], typing.Any] +class _KeyFnUnset: + def __repr__(self) -> str: + return "" + + +_KEY_FN_UNSET = _KeyFnUnset() + + class _MemoFns(typing.NamedTuple): key_fn: _KeyFn state_fn: _StateFn | None = None -_memo_fns: dict[type, _MemoFns] = {} -_memo_type_identifiers: dict[int, tuple[weakref.ReferenceType[type], str]] = {} +_memo_fns: dict[int, tuple[weakref.ReferenceType[type], _MemoFns]] = {} +_stable_type_ids: dict[int, tuple[weakref.ReferenceType[type], str]] = {} class StateFnEntry(typing.NamedTuple): @@ -55,6 +61,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], @@ -123,191 +147,198 @@ def canonical_module_name(obj: typing.Any) -> str: def _memo_type_label(typ: type) -> str: - """Return a user-facing label for type-identifier validation errors.""" + """Return a user-facing label for stable type ID validation errors.""" return f"{canonical_module_name(typ)}.{getattr(typ, '__qualname__', '')}" -def _validate_memo_type_identifier(identifier: object, *, source: str) -> str: - """Validate a non-empty stable memo type identifier.""" - if not isinstance(identifier, str): - raise TypeError(f"{source} must be a str, got {type(identifier).__name__}") - if identifier.strip() == "": +def _validate_stable_type_id(stable_type_id: object, *, source: str) -> str: + """Validate a non-empty stable type ID.""" + if not isinstance(stable_type_id, str): + raise TypeError(f"{source} must be a str, got {type(stable_type_id).__name__}") + if stable_type_id.strip() == "": raise ValueError( f"{source} must be non-empty and contain non-whitespace characters" ) - return identifier + return stable_type_id -def _remove_memo_type_identifier( - type_id: int, dead_ref: weakref.ReferenceType[type] -) -> None: +def _remove_stable_type_id(type_id: int, dead_ref: weakref.ReferenceType[type]) -> None: """Remove ``type_id`` only if ``dead_ref`` is still the stored weakref.""" - entry = _memo_type_identifiers.get(type_id) + entry = _stable_type_ids.get(type_id) if entry is not None and entry[0] is dead_ref: - _memo_type_identifiers.pop(type_id, None) + _stable_type_ids.pop(type_id, None) -def register_memo_type_identifier(typ: type, identifier: str) -> None: - """Register a stable memo type identity for one exact Python type. - - Type-aware fingerprints use it instead of module+qualname. Registration - overrides ``typ.__coco_memo_type_id__``; explicit class-object hooks and - pickle fallback must carry their own stable key. - """ - if not isinstance(typ, type): - raise TypeError( - "register_memo_type_identifier() expects typ to be a type, " - f"got {type(typ).__name__}" - ) - identifier = _validate_memo_type_identifier( - identifier, source="register_memo_type_identifier(..., identifier)" - ) +def _register_stable_type_id(typ: type, stable_type_id: str) -> None: + """Register a stable type ID for one exact Python type.""" type_id = id(typ) - def _remove_stale_type_identifier( + def _remove_stale_type_id( dead_ref: weakref.ReferenceType[type], ) -> None: - _remove_memo_type_identifier(type_id, dead_ref) + _remove_stable_type_id(type_id, dead_ref) - _memo_type_identifiers[type_id] = ( - weakref.ref(typ, _remove_stale_type_identifier), - identifier, + _stable_type_ids[type_id] = ( + weakref.ref(typ, _remove_stale_type_id), + stable_type_id, ) -def _unregister_memo_type_identifier(typ: type) -> None: - """Best-effort test helper for removing an exact-type registration.""" +def _unregister_stable_type_id(typ: type) -> None: + """Best-effort removal of an exact-type stable type ID registration.""" type_id = id(typ) - entry = _memo_type_identifiers.get(type_id) + entry = _stable_type_ids.get(type_id) if entry is not None and entry[0]() is typ: - _memo_type_identifiers.pop(type_id, None) + _stable_type_ids.pop(type_id, None) -def _registered_memo_type_identifier(typ: type) -> str | None: - """Return the registered identifier for ``typ`` from the id-keyed table.""" +def _registered_stable_type_id(typ: type) -> str | None: + """Return the registered stable type ID for ``typ`` from the id-keyed table.""" type_id = id(typ) - entry = _memo_type_identifiers.get(type_id) + entry = _stable_type_ids.get(type_id) if entry is None: return None - ref, identifier = entry + ref, stable_type_id = entry if ref() is typ: - return identifier - _memo_type_identifiers.pop(type_id, None) + return stable_type_id + _stable_type_ids.pop(type_id, None) return None -def _lookup_memo_type_identifier(typ: type) -> str | None: - """Resolve a registered or exact ``__coco_memo_type_id__`` identifier.""" - identifier = _registered_memo_type_identifier(typ) - if identifier is not None: - return identifier +def _lookup_stable_type_id(typ: type) -> str | None: + """Resolve a registered or exact ``__coco_memo_type_id__`` stable type ID.""" + stable_type_id = _registered_stable_type_id(typ) + if stable_type_id is not None: + return stable_type_id if "__coco_memo_type_id__" in typ.__dict__: - return _validate_memo_type_identifier( + return _validate_stable_type_id( typ.__dict__["__coco_memo_type_id__"], source=f"{_memo_type_label(typ)}.__coco_memo_type_id__", ) return None -_MEMO_KEY_ATTR = "__coco_memo_key__" -_MEMO_STATE_ATTR = "__coco_memo_state__" +def _remove_memo_fns(type_id: int, dead_ref: weakref.ReferenceType[type]) -> None: + """Remove ``type_id`` only if ``dead_ref`` is still the stored weakref.""" + entry = _memo_fns.get(type_id) + if entry is not None and entry[0] is dead_ref: + _memo_fns.pop(type_id, None) -def _callable_memo_hook(hook: object) -> _BoundKeyFn | None: - """Return ``hook`` when it is callable.""" - if not callable(hook): - return None - return typing.cast(_BoundKeyFn, hook) +def _register_memo_fns(typ: type, memo_fns: _MemoFns) -> None: + """Register memo functions for one exact Python type.""" + type_id = id(typ) + def _remove_stale_memo_fns(dead_ref: weakref.ReferenceType[type]) -> None: + _remove_memo_fns(type_id, dead_ref) -def _python_callable_accepts_no_args(hook: object) -> bool: - """Return whether a Python function/method can bind zero arguments. + _memo_fns[type_id] = (weakref.ref(typ, _remove_stale_memo_fns), memo_fns) - This is deliberately limited to regular Python callables. Some valid - callable objects and C-extension functions cannot expose an inspectable - signature; those are accepted and validated by the actual call. - """ - if not (inspect.isfunction(hook) or inspect.ismethod(hook)): - return True - try: - inspect.signature(hook).bind() - except TypeError: - return False - return True - - -def _metaclass_object_memo_hook(cls: type) -> _BoundKeyFn | None: - """Resolve a memo key hook from ``cls``'s metaclass.""" - for base in typing.cast(type, type(cls)).__mro__: - raw = base.__dict__.get(_MEMO_KEY_ATTR) - if raw is None: - continue - hook = ( - typing.cast(typing.Any, raw).__get__(cls, type(cls)) - if hasattr(raw, "__get__") - else raw - ) - return _callable_memo_hook(hook) - return None +def _unregister_memo_fns(typ: type) -> None: + """Best-effort removal of an exact-type memo function registration.""" + type_id = id(typ) + entry = _memo_fns.get(type_id) + if entry is not None and entry[0]() is typ: + _memo_fns.pop(type_id, None) -def _class_object_memo_hook(cls: type) -> _BoundKeyFn | None: - """Resolve an explicit memo key hook for a class object. - Plain class-body methods that require an instance are ignored because there - is no ``self`` to bind. Zero-argument functions assigned as class - attributes, descriptor-based hooks, callable objects, and metaclass hooks - remain valid class-object hooks. - """ - for base in cls.__mro__: - raw = base.__dict__.get(_MEMO_KEY_ATTR) - if raw is None: - continue - hook = getattr(cls, _MEMO_KEY_ATTR, None) - if isinstance(raw, (classmethod, staticmethod)): - if not callable(hook) or not _python_callable_accepts_no_args(hook): - hook_kind = ( - "classmethod" if isinstance(raw, classmethod) else "staticmethod" - ) - raise TypeError( - f"{_memo_type_label(base)}.{_MEMO_KEY_ATTR} is a {hook_kind} " - "that cannot be called with zero arguments; class-object hooks " - "must take no arguments after binding" - ) - return typing.cast(_BoundKeyFn, hook) - if callable(hook) and _python_callable_accepts_no_args(hook): - return typing.cast(_BoundKeyFn, hook) - break - return _metaclass_object_memo_hook(cls) - - -def _class_object_state_fn_entry(cls: type) -> StateFnEntry | None: - """Resolve a memo state hook bound to ``cls``'s metaclass.""" - typ = type(cls) - for base in typing.cast(type, typ).__mro__: - raw = base.__dict__.get(_MEMO_STATE_ATTR) - if raw is None: - continue - state_hook = ( - typing.cast(typing.Any, raw).__get__(cls, typ) - if hasattr(raw, "__get__") - else raw - ) - if not callable(state_hook): - return None - raw_fn = getattr(typ, _MEMO_STATE_ATTR) - return _make_state_fn_entry(state_hook, raw_fn) +def _registered_memo_fns(typ: type) -> _MemoFns | None: + """Return registered memo functions for ``typ`` from the id-keyed table.""" + type_id = id(typ) + entry = _memo_fns.get(type_id) + if entry is None: + return None + ref, memo_fns = entry + if ref() is typ: + return memo_fns + _memo_fns.pop(type_id, None) return None +_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) -> tuple[Fingerprintable, Fingerprintable]: - """Return stable-ID or module+qualname type identity parts.""" - identifier = _lookup_memo_type_identifier(typ) - if identifier is not None: - return (("__coco_memo_type_id__", identifier), None) + """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. + """ + stable_type_id = _lookup_stable_type_id(typ) + if stable_type_id is not None: + return (("__coco_memo_type_id__", stable_type_id), None) return (canonical_module_name(typ), getattr(typ, "__qualname__", None)) +def _canonicalize_key_fragment( + obj: object, + state: _CanonicalizeState, + state_methods: list[StateFnEntry], +) -> Fingerprintable: + """Canonicalize a memo-key fragment within the current root traversal. + + Sharing traversal state preserves cycles through the parent object and keeps + temporary fragment objects alive so their IDs cannot be reused during this + traversal. + """ + + return _canonicalize(obj, state, state_methods) + + +def _canonicalize_registered_memo_key( + obj: object, + owner: type, + memo: _MemoFns, + state: _CanonicalizeState, + state_methods: list[StateFnEntry], +) -> Fingerprintable: + key = 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, + *_type_identity_parts(owner), + _canonicalize_key_fragment(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) + for owner in metaclass.__mro__: + if owner is object: + break + memo = _registered_memo_fns(owner) + if memo is not None: + return _canonicalize_registered_memo_key( + cls, owner, memo, state, state_methods + ) + + 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)), + ) + + 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) @@ -320,7 +351,7 @@ def _is_pydantic_model(obj: object) -> bool: def _canonicalize_dataclass( obj: object, - _seen: dict[int, int], + state: _CanonicalizeState, state_methods: list[StateFnEntry], ) -> Fingerprintable: """Canonicalize a dataclass instance. @@ -334,7 +365,7 @@ def _canonicalize_dataclass( "dataclass", *_type_identity_parts(typ), 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 ), ) @@ -342,7 +373,7 @@ def _canonicalize_dataclass( def _canonicalize_pydantic( obj: object, - _seen: dict[int, int], + state: _CanonicalizeState, state_methods: list[StateFnEntry], ) -> Fingerprintable: """Canonicalize a Pydantic v2 model instance. @@ -356,7 +387,7 @@ def _canonicalize_pydantic( "pydantic", *_type_identity_parts(typ), tuple( - (name, _canonicalize(getattr(obj, name), _seen, state_methods)) + (name, _canonicalize(getattr(obj, name), state, state_methods)) for name in field_names ), ) @@ -381,46 +412,124 @@ def __coco_memo_key__(self) -> typing.NoReturn: ) +@typing.overload def register_memo_key_function( - typ: type, key_fn: _KeyFn, *, state_fn: _StateFn | None = None -) -> None: - """Register a memo key function for a type. + typ: type, + key_fn: _KeyFn, + *, + state_fn: _StateFn | None = None, + stable_type_id: str | None = None, +) -> 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``). +@typing.overload +def register_memo_key_function( + typ: type, + *, + stable_type_id: str, +) -> None: ... + + +def register_memo_key_function( + typ: type, + key_fn: _KeyFn | object = _KEY_FN_UNSET, + *, + state_fn: _StateFn | None = None, + stable_type_id: str | None = None, +) -> None: + """Register a memo key function and/or stable type ID for a type. + + Key-function resolution is MRO-aware: the most specific registered base + type wins. Stable type IDs registered without a key function apply to the + exact type only; stable type IDs registered with a key function identify + that selected owner type. Each call replaces the full registration for + ``typ``: omitting ``stable_type_id`` clears any previous stable type ID registered for ``typ``, + and omitting ``key_fn`` clears any previous key/state functions. + + To register only a stable type ID, omit ``key_fn`` rather than passing + ``None``. 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( + "register_memo_key_function() expects typ to be a type, " + f"got {type(typ).__name__}" + ) + if stable_type_id is not None: + stable_type_id = _validate_stable_type_id( + stable_type_id, + source="register_memo_key_function(..., stable_type_id)", + ) + if key_fn is None: + raise TypeError( + "register_memo_key_function() key_fn must be callable; omit key_fn " + "when registering only a stable type ID" + ) + if key_fn is _KEY_FN_UNSET: + if state_fn is not None: + raise TypeError( + "register_memo_key_function() state_fn requires a memo key function" + ) + if stable_type_id is None: + raise TypeError( + "register_memo_key_function() requires a key_fn or stable_type_id" + ) + elif not callable(key_fn): + raise TypeError( + "register_memo_key_function() key_fn must be callable, " + f"got {type(key_fn).__name__}" + ) + if state_fn is not None and not callable(state_fn): + raise TypeError( + "register_memo_key_function() state_fn must be callable, " + f"got {type(state_fn).__name__}" + ) + + if stable_type_id is not None: + _register_stable_type_id(typ, stable_type_id) + else: + _unregister_stable_type_id(typ) + if key_fn is not _KEY_FN_UNSET: + _register_memo_fns(typ, _MemoFns(typing.cast(_KeyFn, key_fn), state_fn)) + else: + _unregister_memo_fns(typ) 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`. - - Example: - import cocoindex as coco - from some_library import StatefulGenerator - - coco.register_not_memo_keyable(StatefulGenerator) + 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) + _unregister_stable_type_id(typ) + _register_memo_fns(typ, _MemoFns(_raise_not_memo_keyable)) def unregister_memo_key_function(typ: type) -> None: - """Remove a previously registered memo key function (best-effort).""" + """Remove registered memo key function and stable type ID (best-effort).""" - _memo_fns.pop(typ, None) + if isinstance(typ, type): + _unregister_stable_type_id(typ) + _unregister_memo_fns(typ) def _stable_sort_key(v: Fingerprintable) -> tuple[typing.Any, ...]: @@ -455,12 +564,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: @@ -471,87 +579,71 @@ def _canonicalize( if isinstance(obj, (bytearray, memoryview)): return bytes(obj) - # 2) Hook / registry (apply once, then recurse on returned key fragment) - hook = ( - _class_object_memo_hook(obj) - if isinstance(obj, type) - else getattr(obj, _MEMO_KEY_ATTR, 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) + + hook = getattr(obj, _MEMO_KEY_ATTR, None) if hook is not None and callable(hook): k = hook() typ = type(obj) tag = "hook" - if isinstance(obj, type): - state_entry = _class_object_state_fn_entry(obj) - if state_entry is not None: - tag = "shook" - state_methods.append(state_entry) - else: - 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, _MEMO_STATE_ATTR) - state_methods.append(_make_state_fn_entry(state_hook, raw_fn)) + 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, _MEMO_STATE_ATTR) + state_methods.append(_make_state_fn_entry(state_hook, raw_fn)) return ( tag, *_type_identity_parts(typ), - _canonicalize(k, _seen, state_methods), + _canonicalize_key_fragment(k, state, state_methods), ) - for base in type(obj).__mro__: - memo = _memo_fns.get(base) + for owner in type(obj).__mro__: + memo = _registered_memo_fns(owner) 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, - *_type_identity_parts(base), - _canonicalize(k, _seen, state_methods), + return _canonicalize_registered_memo_key( + obj, owner, memo, 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, state, state_methods) # 6) Pydantic v2 models if _is_pydantic_model(obj): - return _canonicalize_pydantic(obj, _seen, state_methods) + return _canonicalize_pydantic(obj, state, state_methods) # 7) Fallback try: @@ -579,13 +671,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 ( @@ -601,7 +693,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=[]) ) @@ -636,14 +728,9 @@ def fingerprint_call( return core.fingerprint_simple_object(call_key_obj) -# Register memo key for class types. -register_memo_key_function(type, lambda cls: _type_identity_parts(cls)) - - __all__ = [ "NotMemoKeyable", "register_memo_key_function", - "register_memo_type_identifier", "register_not_memo_keyable", "unregister_memo_key_function", "fingerprint_call", diff --git a/python/tests/core/test_function_memo.py b/python/tests/core/test_function_memo.py index e56541929..6c68b95d8 100644 --- a/python/tests/core/test_function_memo.py +++ b/python/tests/core/test_function_memo.py @@ -88,14 +88,14 @@ def _process_stable_type_source_data() -> None: coco.declare_target_state(GlobalDictTarget.target_state(key, transformed_value)) -def test_stable_memo_type_id_reuses_memo_across_renamed_dataclass() -> None: +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_memo_type_id_reuses_memo_across_renamed_dataclass", + name="test_stable_type_id_reuses_memo_across_renamed_dataclass", environment=coco_env, ), _process_stable_type_source_data, diff --git a/python/tests/core/test_memo_state_validation.py b/python/tests/core/test_memo_state_validation.py index 7e47c6be4..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, @@ -716,7 +717,7 @@ def test_state_changed_but_reusable_component_sync() -> None: # ============================================================================ -# Raw class-object metaclass state validation (sync) +# Raw class-object metaclass memo methods are ignored (sync) # ============================================================================ _class_object_source: dict[str, type] = {} @@ -725,24 +726,30 @@ def test_state_changed_but_reusable_component_sync() -> None: class _ClassObjectMemoMeta(type): def __coco_memo_key__(cls) -> object: - return getattr(cls, "stable_key") + raise AssertionError("class objects must not call metaclass memo key") def __coco_memo_state__(cls, prev_state: Any) -> coco.MemoStateOutcome: - state_value = getattr(cls, "state_value") _class_object_prev_states.append(prev_state) - memo_valid = not coco.is_non_existence(prev_state) and prev_state == state_value - return coco.MemoStateOutcome(state=state_value, memo_valid=memo_valid) - - -def _make_class_object(stable_key: str, state_value: int, content: str) -> type: + 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}", (), - { - "stable_key": stable_key, - "state_value": state_value, - "content": content, - }, + attrs, ) @@ -759,8 +766,8 @@ def _process_class_objects() -> None: coco.declare_target_state(GlobalDictTarget.target_state(key, result)) -def test_class_object_metaclass_state_validation_sync() -> None: - """Class-object metaclass state participates in App memo validation.""" +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() @@ -768,53 +775,97 @@ def test_class_object_metaclass_state_validation_sync() -> None: app = coco.App( coco.AppConfig( - name="test_class_object_metaclass_state_validation_sync", + name="test_class_object_metaclass_memo_methods_are_ignored_sync", environment=coco_env, ), _process_class_objects, ) - # Run 1: no previous state, so the transform executes and persists content A. - _class_object_source["row"] = _make_class_object("stable", 1, "A") + 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 len(_class_object_prev_states) == 1 - assert coco.is_non_existence(_class_object_prev_states[0]) + assert _class_object_prev_states == [] assert GlobalDictTarget.store.data == { "row": DictDataWithPrev( data="class content: A", prev=[], prev_may_be_missing=True ) } - # Run 2: equivalent replacement with same state is a cache hit; C is ignored. - _class_object_source["row"] = _make_class_object("stable", 1, "C") + # 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 len(_class_object_prev_states) == 2 - assert coco.is_non_existence(_class_object_prev_states[0]) - assert _class_object_prev_states[1:] == [1] + assert _class_object_prev_states == [] assert GlobalDictTarget.store.data["row"].data == "class content: A" - # Run 3: same key but new state invalidates the cache and persists content B. - _class_object_source["row"] = _make_class_object("stable", 2, "B") + # 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 len(_class_object_prev_states) == 3 - assert coco.is_non_existence(_class_object_prev_states[0]) - assert _class_object_prev_states[1:] == [1, 1] + assert _class_object_prev_states == [] assert GlobalDictTarget.store.data == { "row": DictDataWithPrev( - data="class content: B", + data="class content: C", prev=["class content: A"], prev_may_be_missing=False, ) } - # Run 4: state 2 was persisted, so another replacement hits cache; D is ignored. - _class_object_source["row"] = _make_class_object("stable", 2, "D") - app.update_blocking() - assert _metrics.collect() == {} - assert len(_class_object_prev_states) == 4 - assert coco.is_non_existence(_class_object_prev_states[0]) - assert _class_object_prev_states[1:] == [1, 1, 2] - assert GlobalDictTarget.store.data["row"].data == "class content: B" + +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 0cb264236..e89bb4e96 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -1,16 +1,16 @@ import dataclasses import math +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, - StateFnEntry, - _unregister_memo_type_identifier, - register_memo_type_identifier, register_memo_key_function, + register_not_memo_keyable, unregister_memo_key_function, ) from cocoindex._internal.typing import MemoStateOutcome @@ -24,6 +24,14 @@ def _dummy_fn(*args: Any, **kwargs: Any) -> None: raise RuntimeError("not called") +def _canonical_contains(value: object, needle: object) -> bool: + if value == needle: + return True + if isinstance(value, tuple): + return any(_canonical_contains(item, needle) for item in value) + return False + + def test_fingerprint_dict_order_independent() -> None: a = {"x": 1, "y": 2} b = {"y": 2, "x": 1} @@ -84,8 +92,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"),), {}, []) @@ -124,9 +132,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),), {}, [] ) @@ -185,30 +193,27 @@ class NewModel(BaseModel): ) != fingerprint_call(_dummy_fn, (NewModel(value=2),), {}, []) -class _UninspectableZeroArgMemoHook: - @property - def __signature__(self) -> object: - raise ValueError("signature unavailable") - - def __call__(self) -> object: - return ("class-hook", "stable") - - -def test_raw_class_object_honors_zero_arg_memo_key() -> None: - def stable_key() -> object: - return ("class-hook", "stable") - - def other_key() -> object: - return ("class-hook", "other") - +def test_raw_class_object_stable_type_id_never_calls_staticmethod_memo_key() -> None: class OldEntry: - __coco_memo_key__ = stable_key + __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_key__ = stable_key + __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_key__ = other_key + __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" @@ -217,223 +222,409 @@ class ChangedEntry: assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( _dummy_fn, (NewEntry,), {}, [] ) - assert fingerprint_call(_dummy_fn, (NewEntry,), {}, []) != fingerprint_call( + assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) != fingerprint_call( _dummy_fn, (ChangedEntry,), {}, [] ) -def test_raw_class_object_metaclass_hook_collects_state_method() -> None: - class MemoMeta(type): - def __coco_memo_key__(cls) -> object: - return ("metaclass-hook", "stable") - - def __coco_memo_state__(cls, prev_state: object) -> MemoStateOutcome: - return MemoStateOutcome( - state=("state", cls.__name__, prev_state), memo_valid=True - ) +def test_registered_stable_type_id_applies_to_class_objects() -> None: + class OldEntry: + pass - class OldEntry(metaclass=MemoMeta): + class NewEntry: pass - class NewEntry(metaclass=MemoMeta): + class ChangedEntry: pass - OldEntry.__module__ = "tests.old_raw_class" - NewEntry.__module__ = "tests.new_raw_class" + try: + register_memo_key_function( + OldEntry, stable_type_id="test.RegisteredRawClass/v1" + ) + register_memo_key_function( + NewEntry, stable_type_id="test.RegisteredRawClass/v1" + ) + register_memo_key_function( + ChangedEntry, stable_type_id="test.RegisteredRawClass/v2" + ) - old_methods: list[Any] = [] - new_methods: list[Any] = [] - assert fingerprint_call( - _dummy_fn, (OldEntry,), {}, old_methods - ) == fingerprint_call(_dummy_fn, (NewEntry,), {}, new_methods) - - assert len(old_methods) == 1 - assert len(new_methods) == 1 - assert isinstance(old_methods[0], StateFnEntry) - assert isinstance(new_methods[0], StateFnEntry) - assert old_methods[0].call("previous") == MemoStateOutcome( - state=("state", "OldEntry", "previous"), memo_valid=True - ) - assert new_methods[0].call("previous") == MemoStateOutcome( - state=("state", "NewEntry", "previous"), memo_valid=True - ) + 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_key_function(OldEntry) + unregister_memo_key_function(NewEntry) + unregister_memo_key_function(ChangedEntry) -def test_raw_class_object_accepts_uninspectable_zero_arg_staticmethod_memo_key() -> ( - None -): - hook = _UninspectableZeroArgMemoHook() +def test_hook_memo_key_fragment_preserves_parent_cycle() -> None: + class Entry: + def __init__(self, parent: list[object]) -> None: + self.parent = parent - class OldEntry: - __coco_memo_key__ = staticmethod(hook) + def __coco_memo_key__(self) -> object: + return self.parent - class NewEntry: - __coco_memo_key__ = staticmethod(hook) + def make_graph() -> list[object]: + parent: list[object] = [] + parent.append(Entry(parent)) + return parent - OldEntry.__module__ = "tests.old_raw_class" - NewEntry.__module__ = "tests.new_raw_class" + graph = make_graph() + canonical = _memo_fingerprint._canonicalize(graph, None, []) - assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( - _dummy_fn, (NewEntry,), {}, [] + assert canonical == ( + "seq", + (("hook", *_memo_fingerprint._type_identity_parts(Entry), ("ref", 0)),), ) + assert _memo_fingerprint.memo_fingerprint( + graph + ) == _memo_fingerprint.memo_fingerprint(make_graph()) -def test_raw_class_object_accepts_uninspectable_zero_arg_plain_memo_key() -> None: - hook = _UninspectableZeroArgMemoHook() +def test_registered_memo_key_fragment_preserves_parent_cycle() -> None: + class Entry: + def __init__(self, parent: list[object]) -> None: + self.parent = parent - class OldEntry: - __coco_memo_key__ = hook + def make_graph() -> list[object]: + parent: list[object] = [] + parent.append(Entry(parent)) + return parent - class NewEntry: - __coco_memo_key__ = hook + try: + register_memo_key_function(Entry, lambda entry: entry.parent) + graph = make_graph() + canonical = _memo_fingerprint._canonicalize(graph, None, []) - OldEntry.__module__ = "tests.old_raw_class" - NewEntry.__module__ = "tests.new_raw_class" + assert canonical == ( + "seq", + (("hook", *_memo_fingerprint._type_identity_parts(Entry), ("ref", 0)),), + ) + assert _memo_fingerprint.memo_fingerprint( + graph + ) == _memo_fingerprint.memo_fingerprint(make_graph()) + finally: + unregister_memo_key_function(Entry) - assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( - _dummy_fn, (NewEntry,), {}, [] - ) +def test_hook_memo_key_fragments_remain_alive_for_root_traversal() -> None: + class Fragment(list[object]): + pass -def test_raw_class_object_rejects_non_zero_arg_staticmethod_memo_key() -> None: - """Reject invalid staticmethod hooks before calling class-object memo keys.""" + first_fragment_ref: weakref.ReferenceType[Fragment] | None = None - class Entry: - @staticmethod - def __coco_memo_key__(value: object) -> object: - return ("class-hook", value) - - with pytest.raises( - TypeError, - match=( - r"Entry\.__coco_memo_key__ is a staticmethod that cannot be called " - r"with zero arguments; class-object hooks must take no arguments " - r"after binding" - ), - ): - fingerprint_call(_dummy_fn, (Entry,), {}, []) + 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), + ("seq", ("first",)), + ), + ( + "hook", + *_memo_fingerprint._type_identity_parts(SecondEntry), + ("seq", ("second",)), + ), + ), + ) -def test_raw_class_object_rejects_non_zero_arg_classmethod_memo_key() -> None: - """Reject invalid classmethod hooks after descriptor binding leaves arguments.""" +def test_memo_key_fragment_preserves_shared_reference_ordinals() -> None: class Entry: - @classmethod - def __coco_memo_key__(cls, value: object) -> object: - return ("class-hook", cls.__name__, value) - - with pytest.raises( - TypeError, - match=( - r"Entry\.__coco_memo_key__ is a classmethod that cannot be called " - r"with zero arguments; class-object hooks must take no arguments " - r"after binding" + 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), + ("seq", (("seq", ("shared",)), ("ref", 1))), + ) + assert parent_wrapped == ( + "seq", + ( + ( + "hook", + *_memo_fingerprint._type_identity_parts(Entry), + ("seq", (("seq", ("shared",)), ("ref", 2))), + ), ), - ): - fingerprint_call(_dummy_fn, (Entry,), {}, []) + ) -def test_raw_class_object_ignores_instance_memo_state() -> None: +def test_raw_class_object_default_identity_never_calls_classmethod_memo_key() -> None: class Entry: - @staticmethod - def __coco_memo_key__() -> object: - return ("class-hook", "stable") + @classmethod + def __coco_memo_key__(cls) -> object: + raise AssertionError("class-object fingerprint must not call memo key") - def __coco_memo_state__(self, prev_state: object) -> MemoStateOutcome: - raise AssertionError("instance memo state must not run for class objects") + class OtherEntry: + @classmethod + def __coco_memo_key__(cls) -> object: + raise AssertionError("class-object fingerprint must not call memo key") - methods: list[Any] = [] - fp1 = fingerprint_call(_dummy_fn, (Entry,), {}, methods) - fp2 = fingerprint_call(_dummy_fn, (Entry,), {}, []) + 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,), {}, [] + ) - assert fp1 == fp2 - assert methods == [] +def test_raw_class_object_honors_registered_metaclass_memo_key_and_state() -> None: + key_calls: list[type] = [] + state_calls: list[tuple[type, object]] = [] -def test_raw_class_object_uses_metaclass_hook_after_ignored_instance_key() -> None: class MemoMeta(type): def __coco_memo_key__(cls) -> object: - return ("metaclass-hook", "stable") + raise AssertionError("raw classes must not call metaclass memo attributes") - class Base: - def __coco_memo_key__(self) -> object: - raise AssertionError("instance memo key must not run for class objects") + def __coco_memo_state__(cls, prev_state: object) -> MemoStateOutcome: + raise AssertionError("raw classes must not call metaclass memo attributes") - class OldEntry(Base, metaclass=MemoMeta): + class OldEntry(metaclass=MemoMeta): pass - class NewEntry(Base, metaclass=MemoMeta): + class NewEntry(metaclass=MemoMeta): pass - OldEntry.__module__ = "tests.old_raw_class" - NewEntry.__module__ = "tests.new_raw_class" + def metaclass_key(cls: type) -> object: + key_calls.append(cls) + return ("metaclass", cls.__name__) - assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( - _dummy_fn, (NewEntry,), {}, [] - ) + 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_key_function( + 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 + ) -def test_raw_class_object_stable_type_id_ignores_instance_memo_key() -> None: - class OldEntry: - __coco_memo_type_id__ = "test.RawClass/v1" + 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_key_function(MemoMeta) - def __coco_memo_key__(self) -> object: - raise AssertionError("instance memo key must not run for class objects") - class NewEntry: - __coco_memo_type_id__ = "test.RawClass/v1" +def test_raw_class_object_honors_registered_type_memo_key_and_state() -> None: + key_calls: list[type] = [] + state_calls: list[tuple[type, object]] = [] - def __coco_memo_key__(self) -> object: - raise AssertionError("instance memo key must not run for class objects") + class MemoAttributesMustNotRun: + @classmethod + def __coco_memo_key__(cls) -> object: + raise AssertionError("raw classes must not call class memo attributes") - class ChangedEntry: - __coco_memo_type_id__ = "test.RawClass/v2" + @classmethod + def __coco_memo_state__(cls, prev_state: object) -> MemoStateOutcome: + raise AssertionError("raw classes must not call class memo attributes") - def __coco_memo_key__(self) -> object: - raise AssertionError("instance memo key must not run for class objects") + class OldEntry(MemoAttributesMustNotRun): + pass - OldEntry.__module__ = "tests.old_raw_class" - NewEntry.__module__ = "tests.new_raw_class" - ChangedEntry.__module__ = "tests.new_raw_class" + class NewEntry(MemoAttributesMustNotRun): + pass - assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( - _dummy_fn, (NewEntry,), {}, [] - ) - assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) != fingerprint_call( - _dummy_fn, (ChangedEntry,), {}, [] - ) + 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" + ) -def test_raw_class_object_explicit_hook_owns_stable_namespace() -> None: - """Explicit class-object hooks are custom keys, not stable-ID wrappers.""" + stable_type_id = "test.RawClassRegisteredTypeOwner/v1" + try: + register_memo_key_function( + 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 + ) - def shared_key() -> object: - return ("class-hook", "shared") + 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_key_function(type) - def changed_key() -> object: - return ("class-hook", "changed") - class OldEntry: - __coco_memo_type_id__ = "test.RawClassHook/v1" - __coco_memo_key__ = staticmethod(shared_key) +def test_raw_class_object_ignores_registered_object_memo_key() -> None: + object_key_calls: list[object] = [] - class ChangedKeyEntry: - __coco_memo_type_id__ = "test.RawClassHook/v1" - __coco_memo_key__ = staticmethod(changed_key) + class Entry: + pass - class ChangedTypeIdEntry: - __coco_memo_type_id__ = "test.RawClassHook/v2" - __coco_memo_key__ = staticmethod(shared_key) + expected = fingerprint_call(_dummy_fn, (Entry,), {}, []) - assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) != fingerprint_call( - _dummy_fn, (ChangedKeyEntry,), {}, [] - ) - assert fingerprint_call(_dummy_fn, (OldEntry,), {}, []) == fingerprint_call( - _dummy_fn, (ChangedTypeIdEntry,), {}, [] + def object_key(obj: object) -> object: + object_key_calls.append(obj) + return "object memo key ran" + + try: + register_memo_key_function(object, object_key) + assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == expected + assert object_key_calls == [] + finally: + unregister_memo_key_function(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_key_function(type, stable_type_id="test.RegisteredType/v1") + assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == original + finally: + unregister_memo_key_function(type) + + +def test_raw_class_object_stable_type_id_is_exact_type() -> None: + class Parent: + __coco_memo_type_id__ = "test.RawClassExact/v1" + + class Child(Parent): + pass + + assert fingerprint_call(_dummy_fn, (Parent,), {}, []) != fingerprint_call( + _dummy_fn, (Child,), {}, [] ) -def test_registered_memo_type_identifier_allows_renamed_type_reuse() -> None: +def test_register_memo_key_function_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_key_function( + 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_key_function(Base) + + +def test_register_memo_key_function_registers_stable_type_id_without_key_function() -> ( + None +): class OldEntry: def __init__(self, value: object) -> None: self.value = value @@ -448,9 +639,9 @@ def __init__(self, value: object) -> None: def __coco_memo_key__(self) -> object: return ("entry", self.value) - register_memo_type_identifier(OldEntry, "test.RegisteredEntry/v1") - register_memo_type_identifier(NewEntry, "test.RegisteredEntry/v1") try: + register_memo_key_function(OldEntry, stable_type_id="test.RegisteredEntry/v1") + register_memo_key_function(NewEntry, stable_type_id="test.RegisteredEntry/v1") assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( _dummy_fn, (NewEntry(1),), {}, [] ) @@ -458,38 +649,407 @@ def __coco_memo_key__(self) -> object: _dummy_fn, (NewEntry(2),), {}, [] ) finally: - _unregister_memo_type_identifier(OldEntry) - _unregister_memo_type_identifier(NewEntry) + unregister_memo_key_function(OldEntry) + unregister_memo_key_function(NewEntry) -def test_registered_memo_key_function_uses_registered_base_type_identity() -> None: - class Base: +def test_stable_type_id_only_registration_is_exact_for_subclasses() -> None: + class Parent: def __init__(self, value: object) -> None: self.value = value - class ChildA(Base): - pass + def __coco_memo_key__(self) -> object: + return ("entry", self.value) - class ChildB(Base): + class Child(Parent): pass - register_memo_key_function(Base, lambda entry: ("base", entry.value)) - register_memo_type_identifier(Base, "test.RegisteredBase/v1") + class SameStableTypeId: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + try: - assert fingerprint_call(_dummy_fn, (ChildA(1),), {}, []) == fingerprint_call( - _dummy_fn, (ChildB(1),), {}, [] + register_memo_key_function( + Parent, stable_type_id="test.RegisteredExactParent/v1" ) - assert fingerprint_call(_dummy_fn, (ChildA(1),), {}, []) != fingerprint_call( - _dummy_fn, (ChildB(2),), {}, [] + register_memo_key_function( + SameStableTypeId, stable_type_id="test.RegisteredExactParent/v1" + ) + + assert fingerprint_call(_dummy_fn, (Parent(1),), {}, []) == fingerprint_call( + _dummy_fn, (SameStableTypeId(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Parent(1),), {}, []) != fingerprint_call( + _dummy_fn, (Child(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Parent,), {}, []) != fingerprint_call( + _dummy_fn, (Child,), {}, [] ) finally: - unregister_memo_key_function(Base) - _unregister_memo_type_identifier(Base) + unregister_memo_key_function(Parent) + unregister_memo_key_function(SameStableTypeId) + + +def test_stable_type_id_only_registration_replaces_key_and_state_functions() -> None: + @dataclasses.dataclass + class Entry: + value: int + marker: str + + def state_fn(obj: Entry, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome(state=prev_state, memo_valid=True) + + try: + register_memo_key_function( + Entry, + lambda entry: ("constant",), + state_fn=state_fn, + ) + methods: list[Any] = [] + assert fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) == ( + fingerprint_call(_dummy_fn, (Entry(2, "b"),), {}, []) + ) + assert len(methods) == 1 + + register_memo_key_function( + Entry, + stable_type_id="test.ReplaceKeyWithStable/v1", + ) + methods = [] + assert fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) != ( + fingerprint_call(_dummy_fn, (Entry(2, "b"),), {}, []) + ) + assert methods == [] + finally: + unregister_memo_key_function(Entry) + + +def test_key_only_registration_replaces_stable_type_id() -> None: + class Entry: + def __init__(self, value: object) -> None: + self.value = value + + class SameStableTypeId: + def __init__(self, value: object) -> None: + self.value = value + + try: + register_memo_key_function(Entry, stable_type_id="test.ReplaceStableWithKey/v1") + register_memo_key_function( + SameStableTypeId, + lambda entry: ("entry", entry.value), + stable_type_id="test.ReplaceStableWithKey/v1", + ) + + register_memo_key_function(Entry, lambda entry: ("entry", entry.value)) + assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) != fingerprint_call( + _dummy_fn, (SameStableTypeId(1),), {}, [] + ) + finally: + unregister_memo_key_function(Entry) + unregister_memo_key_function(SameStableTypeId) + + +def test_key_only_registration_falls_back_to_class_declared_stable_type_id() -> None: + class Entry: + __coco_memo_type_id__ = "test.DeclaredFallback/v1" + + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class DeclaredPeer: + __coco_memo_type_id__ = "test.DeclaredFallback/v1" + + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class RegisteredPeer: + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) == fingerprint_call( + _dummy_fn, (DeclaredPeer(1),), {}, [] + ) + + try: + register_memo_key_function(Entry, stable_type_id="test.RegisteredOverride/v1") + register_memo_key_function( + RegisteredPeer, stable_type_id="test.RegisteredOverride/v1" + ) + + assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) == fingerprint_call( + _dummy_fn, (RegisteredPeer(1),), {}, [] + ) + assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) != fingerprint_call( + _dummy_fn, (DeclaredPeer(1),), {}, [] + ) + + register_memo_key_function(Entry, lambda entry: ("entry", entry.value)) + assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) == fingerprint_call( + _dummy_fn, (DeclaredPeer(1),), {}, [] + ) + finally: + unregister_memo_key_function(Entry) + unregister_memo_key_function(RegisteredPeer) + + +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 + ) -def test_registered_memo_type_identifier_is_identity_exact_for_equal_metaclasses() -> ( + try: + register_memo_key_function( + OldEntry, + key_fn, + state_fn=state_fn, + stable_type_id="test.CombinedStateStable/v1", + ) + register_memo_key_function( + 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_key_function(OldEntry) + unregister_memo_key_function(NewEntry) + + +def test_register_memo_key_function_full_registration_replaces_previous_full_registration() -> ( None ): + class Entry: + def __init__(self, value: object, marker: object) -> None: + self.value = value + self.marker = marker + + def old_state_fn(entry: Any, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome(state=("old", entry.value, prev_state), memo_valid=True) + + def new_state_fn(entry: Any, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome( + state=("new", entry.marker, prev_state), memo_valid=True + ) + + try: + register_memo_key_function( + Entry, + lambda entry: ("old", entry.value), + state_fn=old_state_fn, + stable_type_id="test.ReplaceFull/old", + ) + first_fingerprint = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) + assert first_fingerprint == fingerprint_call( + _dummy_fn, (Entry(1, "b"),), {}, [] + ) + assert first_fingerprint != fingerprint_call( + _dummy_fn, (Entry(2, "a"),), {}, [] + ) + methods: list[Any] = [] + fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) + assert len(methods) == 1 + assert methods[0].call("prev").state == ("old", 1, "prev") + + register_memo_key_function( + Entry, + lambda entry: ("new", entry.marker), + state_fn=new_state_fn, + stable_type_id="test.ReplaceFull/new", + ) + second_fingerprint = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) + assert first_fingerprint != second_fingerprint + assert second_fingerprint == fingerprint_call( + _dummy_fn, (Entry(2, "a"),), {}, [] + ) + assert second_fingerprint != fingerprint_call( + _dummy_fn, (Entry(1, "b"),), {}, [] + ) + methods = [] + fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) + assert len(methods) == 1 + assert methods[0].call("prev").state == ("new", "a", "prev") + finally: + unregister_memo_key_function(Entry) + + +def test_register_not_memo_keyable_replaces_stable_type_id_for_class_objects() -> None: + class Entry: + pass + + class SameStableTypeId: + pass + + try: + register_memo_key_function( + Entry, stable_type_id="test.NotMemoKeyableReplacesStable/v1" + ) + register_memo_key_function( + 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_key_function(Entry) + unregister_memo_key_function(SameStableTypeId) + + +def test_register_memo_key_function_rejects_explicit_none_key_function() -> None: + class Entry: + pass + + with pytest.raises(TypeError, match="key_fn"): + register_memo_key_function( + Entry, cast(Any, None), stable_type_id="test.ExplicitNone/v1" + ) + + +@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", + ), + ], +) +def test_register_memo_key_function_rejects_invalid_forms( + args: tuple[Any, ...], + kwargs: dict[str, Any], + match: str, +) -> None: + class Entry: + pass + + with pytest.raises(TypeError, match=match): + register_memo_key_function(Entry, *args, **kwargs) + + +def test_register_memo_key_function_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_key_function(Entry, **kwargs) + + +def test_unregister_memo_key_function_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_key_function( + RegisteredOnly, + lambda entry: ("registered", entry.value), + stable_type_id="test.UnregisterCombined/v1", + ) + register_memo_key_function( + 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_key_function(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_key_function(RegisteredOnly) + unregister_memo_key_function(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_key_function(OldEntry, stable_type_id="test.UnregisterStable/v1") + register_memo_key_function(NewEntry, stable_type_id="test.UnregisterStable/v1") + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry(1),), {}, [] + ) + unregister_memo_key_function(OldEntry) + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) != fingerprint_call( + _dummy_fn, (NewEntry(1),), {}, [] + ) + finally: + unregister_memo_key_function(OldEntry) + unregister_memo_key_function(NewEntry) + + +def test_registered_stable_type_id_is_identity_exact_for_equal_metaclasses() -> None: class EqMeta(type): def __eq__(cls, other: object) -> bool: return isinstance(other, EqMeta) @@ -509,16 +1069,93 @@ def __coco_memo_key__(self) -> object: _dummy_fn, (B(),), {}, [] ) - register_memo_type_identifier(A, "test.EqualityMetaA/v1") try: + register_memo_key_function(A, stable_type_id="test.EqualityMetaA/v1") assert fingerprint_call(_dummy_fn, (A(),), {}, []) != fingerprint_call( _dummy_fn, (B(),), {}, [] ) finally: - _unregister_memo_type_identifier(A) + unregister_memo_key_function(A) + + +def test_unregister_memo_key_function_handles_unhashable_stable_type_id_only() -> None: + class EqNoHashMeta(type): + def __eq__(cls, other: object) -> bool: + return cls is other + + class OldEntry(metaclass=EqNoHashMeta): + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class NewEntry(metaclass=EqNoHashMeta): + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + try: + register_memo_key_function(OldEntry, stable_type_id="test.UnhashableMeta/v1") + register_memo_key_function(NewEntry, stable_type_id="test.UnhashableMeta/v1") + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( + _dummy_fn, (NewEntry(1),), {}, [] + ) + unregister_memo_key_function(OldEntry) + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) != fingerprint_call( + _dummy_fn, (NewEntry(1),), {}, [] + ) + finally: + unregister_memo_key_function(OldEntry) + unregister_memo_key_function(NewEntry) + + +def test_unregister_memo_key_function_is_identity_exact_for_equal_metaclasses() -> None: + class EqHashMeta(type): + def __eq__(cls, other: object) -> bool: + return isinstance(other, EqHashMeta) + + def __hash__(cls) -> int: + return 1 + + class StableOnly(metaclass=EqHashMeta): + pass + + class Registered(metaclass=EqHashMeta): + def __getstate__(self) -> object: + raise TypeError("registered test object is not picklable") + + try: + register_memo_key_function(Registered, lambda entry: ("registered",)) + register_memo_key_function(StableOnly, stable_type_id="test.EqualUnregister/v1") + unregister_memo_key_function(StableOnly) + fingerprint_call(_dummy_fn, (Registered(),), {}, []) + finally: + unregister_memo_key_function(StableOnly) + unregister_memo_key_function(Registered) def test_stable_type_id_exact_type_and_validation() -> None: + class OldEntry: + __coco_memo_type_id__ = "test.DirectEntry/v1" + + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + + class NewEntry: + __coco_memo_type_id__ = "test.DirectEntry/v1" + + def __init__(self, value: object) -> None: + self.value = value + + def __coco_memo_key__(self) -> object: + return ("entry", self.value) + class Parent: __coco_memo_type_id__ = "test.Parent/v1" @@ -540,6 +1177,12 @@ class EmptyId: def __coco_memo_key__(self) -> object: return ("bad", 1) + 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),), {}, [] + ) assert fingerprint_call(_dummy_fn, (Parent(),), {}, []) != fingerprint_call( _dummy_fn, (Child(),), {}, [] ) @@ -549,21 +1192,27 @@ def __coco_memo_key__(self) -> object: fingerprint_call(_dummy_fn, (EmptyId(),), {}, []) -def test_register_memo_type_identifier_validation_and_export() -> None: +def test_register_memo_key_function_validation_and_public_export() -> None: import cocoindex as coco class Entry: pass - assert coco.register_memo_type_identifier is register_memo_type_identifier + assert coco.register_memo_key_function is register_memo_key_function + assert "register_memo_type_identifier" not in coco.__all__ + assert not hasattr(coco, "register_memo_type_identifier") + assert "register_not_memo_keyable" not in coco.__all__ + assert not hasattr(coco, "register_not_memo_keyable") with pytest.raises(TypeError, match="expects typ to be a type"): - register_memo_type_identifier(cast(Any, object()), "test.Invalid/v1") + register_memo_key_function( + cast(Any, object()), stable_type_id="test.Invalid/v1" + ) with pytest.raises(TypeError, match="must be a str"): - register_memo_type_identifier(Entry, cast(Any, object())) + register_memo_key_function(Entry, stable_type_id=cast(Any, object())) with pytest.raises(ValueError, match="non-empty"): - register_memo_type_identifier(Entry, "") + register_memo_key_function(Entry, stable_type_id="") with pytest.raises(ValueError, match="non-empty"): - register_memo_type_identifier(Entry, " ") + register_memo_key_function(Entry, stable_type_id=" ") def test_cycles_are_supported_and_deterministic() -> None: @@ -977,8 +1626,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 From 3ff15674fbb2282e1c82514f4e4766d3e04cb971 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 18 Jul 2026 12:50:08 -0700 Subject: [PATCH 07/16] feat(memo): add previous type identity support --- .../docs/advanced_topics/memoization_keys.mdx | 39 ++- python/cocoindex/_internal/api.py | 2 + .../cocoindex/_internal/memo_fingerprint.py | 47 +++ python/tests/core/test_function_memo.py | 54 ++++ .../tests/internal/test_memo_fingerprint.py | 302 ++++++++++++++++++ 5 files changed, 436 insertions(+), 8 deletions(-) diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index 850f992ab..a6c223d69 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -28,10 +28,9 @@ Timing depends on where the value comes from. Function arguments are fingerprint For each data value, CocoIndex derives a canonical form with this precedence: -1. **Class objects** (`type`) — the class values themselves, such as `ProductRow`, not instances such as `ProductRow(...)`. CocoIndex first checks for an explicit memo key registration on the class object's metaclass or its bases, including `type`. If none exists, it uses the built-in class-object identity path. It never calls memo-key or memo-state attributes on the class object itself. See the class-object notes below for what is and isn't included. -2. For non-class objects, if the object implements **`__coco_memo_key__()`**, CocoIndex uses its return value. -3. Otherwise, if you registered a **memo key function** for the object's type, CocoIndex uses that. -4. Otherwise, CocoIndex uses built-in canonicalization for supported automatic types, including primitives, containers, dataclasses, Pydantic models, and pickle fallback. +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): @@ -126,11 +125,35 @@ class SourceEntry: content: str ``` -To preserve reuse **after** a refactor, add the stable type ID before moving or renaming the class and ensure every affected memoized call executes under that ID. This seeding execution is cold for affected values: the stable type ID replaces the old module-plus-qualified-name namespace, so CocoIndex creates new memo entries instead of reusing the old ones. +**Choose a durable string ID proactively.** For a new type, or before any memo entries for it matter, use a globally unique durable string `__coco_memo_type_id__` as shown above when memo stability is important and a future move or rename cannot be ruled out. Introducing or changing an ordinary string ID after entries exist creates a new memo namespace and causes one cold execution for affected memoized calls. -An application update seeds only memoized calls that actually execute. A cached parent does not run its body or nested `@coco.fn` calls, so invalidate or reprocess affected parents—or otherwise invoke the affected nested calls—before the refactor. +**Retain an existing automatic identity.** If a type already has reusable entries under its automatic identity, assign that exact previous identity to the moved class: -After seeding, move or rename the class while keeping the same stable type ID so later updates can reuse the entries created under that ID. Stable type IDs preserve an input type's identity across refactors; normal function logic tracking still invalidates memo entries when the decorated function's source changes. Adding the stable type ID only after the move does not reconstruct entries created under the old module plus qualified name. +```python +from dataclasses import dataclass +from typing import ClassVar +import cocoindex as coco + +@dataclass +class SourceEntry: + __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id( + "old_package.models", + "SourceEntry", + ) + + name: str + version: int +``` + +The arguments to `prev_type_id` are the exact previous canonical module and qualified name. Nested classes retain their enclosing names in the qualified name, such as `"Container.SourceEntry"`. A file-backed class previously defined in `__main__` uses the entry script's basename without `.py`; for example, `main.py` becomes `"main"`. + +`prev_type_id` does not scan, copy, promote, or migrate records. For a type-aware canonical form that already included the type's automatic identity, it keeps emitting the same `(old_module, old_qualname)`, so those existing entries remain directly addressable. Entries created through pickle fallback did not include this identity. Adding `prev_type_id` alone has no effect on pickle-fallback canonicalization and cannot recover those entries. Adding a memo key function makes future entries type-aware, changes the canonical form, and therefore starts cold; it still cannot recover prior pickle-fallback entries. Keep `prev_type_id` on the moved class for as long as that historical identity is desired. Replacing it later with an ordinary string stable ID intentionally starts a new namespace. Normal decorated-function logic changes can still invalidate memoization. + +`prev_type_id` preserves only this input type's identity within the memo key. Reuse still requires every other call-key component and normal reuse check to remain compatible; moving or renaming the memoized function can therefore start cold independently of the input type. + +Treat that previous canonical `(module, qualname)` as reserved while any active type emits it, and afterward for as long as entries under the historical identity may still be addressed. If a semantically different type later uses the same module and qualified name, matching memo-key payloads can share cached results. + +Assign the value returned by `prev_type_id` directly, or pass it intact as `stable_type_id`. Converting or serializing it to a plain string removes the historical-identity marker. For text-only persistence, store the module and qualified name separately and call `prev_type_id` again after loading. If you control the class definition, set `__coco_memo_type_id__` on the class. If you cannot or do not want to edit the class definition, register the same stable type ID with `register_memo_key_function`. @@ -266,7 +289,7 @@ coco.register_memo_key_function( # MRO lookup and handles the Child value. ``` -Stable type IDs do not retroactively migrate memo entries created before the stable type ID existed. +Adding a new ordinary string stable type ID does not retroactively migrate memo entries created under the automatic identity or a different string ID. Use `prev_type_id(...)` to retain a previous automatic identity; it keeps those entries directly addressable rather than migrating them. #### Class objects as inputs diff --git a/python/cocoindex/_internal/api.py b/python/cocoindex/_internal/api.py index ca97a3a1a..fe9b6267a 100644 --- a/python/cocoindex/_internal/api.py +++ b/python/cocoindex/_internal/api.py @@ -103,6 +103,7 @@ from .memo_fingerprint import ( memo_fingerprint, + prev_type_id, register_memo_key_function, NotMemoKeyable, ) @@ -943,6 +944,7 @@ class Cursor: "serialize_by_pickle", # .memo_fingerprint "memo_fingerprint", + "prev_type_id", "register_memo_key_function", "NotMemoKeyable", # .pending_marker diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index 4c3c137bc..50adc583f 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -151,6 +151,34 @@ def _memo_type_label(typ: type) -> str: return f"{canonical_module_name(typ)}.{getattr(typ, '__qualname__', '')}" +class _PreviousTypeId(str): + """A prior automatic type identity carried through the stable-ID path.""" + + __slots__ = () + + def __new__(cls, module: str, qualname: str | None = None) -> _PreviousTypeId: + payload = module if qualname is None else f"{len(module)}:{module}{qualname}" + return super().__new__(cls, payload) + + def __getnewargs__(self) -> tuple[str]: + return (str(self),) + + def _identity_parts(self) -> tuple[str, str]: + module_length_str, separator, payload = self.partition(":") + if separator == "": + raise ValueError("invalid previous type identity payload") + module_length = int(module_length_str) + return payload[:module_length], payload[module_length:] + + @property + def module(self) -> str: + return self._identity_parts()[0] + + @property + def qualname(self) -> str: + return self._identity_parts()[1] + + def _validate_stable_type_id(stable_type_id: object, *, source: str) -> str: """Validate a non-empty stable type ID.""" if not isinstance(stable_type_id, str): @@ -162,6 +190,22 @@ def _validate_stable_type_id(stable_type_id: object, *, source: str) -> str: return stable_type_id +def _validate_previous_type_id_part(value: object, *, source: str) -> str: + """Validate and normalize one previous automatic identity part.""" + if isinstance(value, str): + return _validate_stable_type_id(str.__str__(value), source=source) + return _validate_stable_type_id(value, source=source) + + +def prev_type_id(module: str, qualname: str) -> str: + """Return a marker that reuses a type's prior automatic identity.""" + module = _validate_previous_type_id_part(module, source="prev_type_id() module") + qualname = _validate_previous_type_id_part( + qualname, source="prev_type_id() qualname" + ) + return _PreviousTypeId(module, qualname) + + def _remove_stable_type_id(type_id: int, dead_ref: weakref.ReferenceType[type]) -> None: """Remove ``type_id`` only if ``dead_ref`` is still the stored weakref.""" entry = _stable_type_ids.get(type_id) @@ -273,6 +317,8 @@ def _type_identity_parts(typ: type) -> tuple[Fingerprintable, Fingerprintable]: names; ``None`` fills the qualname slot. """ stable_type_id = _lookup_stable_type_id(typ) + if isinstance(stable_type_id, _PreviousTypeId): + return stable_type_id._identity_parts() if stable_type_id is not None: return (("__coco_memo_type_id__", stable_type_id), None) return (canonical_module_name(typ), getattr(typ, "__qualname__", None)) @@ -730,6 +776,7 @@ def fingerprint_call( __all__ = [ "NotMemoKeyable", + "prev_type_id", "register_memo_key_function", "register_not_memo_keyable", "unregister_memo_key_function", diff --git a/python/tests/core/test_function_memo.py b/python/tests/core/test_function_memo.py index 6c68b95d8..fb7eeec6b 100644 --- a/python/tests/core/test_function_memo.py +++ b/python/tests/core/test_function_memo.py @@ -45,6 +45,25 @@ class _StableNewEntry: __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 @@ -123,6 +142,41 @@ def test_stable_type_id_reuses_memo_across_renamed_dataclass() -> None: 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/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index e89bb4e96..b3a56943c 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -1,5 +1,8 @@ +import copy import dataclasses import math +import pickle +import sys import weakref from typing import Any, ClassVar, cast @@ -168,6 +171,210 @@ class Entry: ) +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 + assert len({left, right}) == 2 + + class HostileString(str): + def __len__(self) -> int: + return 1 + + def __str__(self) -> str: + return "wrong-str" + + def __format__(self, format_spec: str) -> str: + return "wrong-format" + + def strip(self, chars: str | None = None) -> str: + raise AssertionError("subclass strip should not be called") + + normalized_marker = cast( + Any, + coco.prev_type_id(HostileString("old.package"), HostileString("Outer.Entry")), + ) + assert (normalized_marker.module, normalized_marker.qualname) == ( + "old.package", + "Outer.Entry", + ) + assert normalized_marker == coco.prev_type_id("old.package", "Outer.Entry") + with pytest.raises(ValueError, match="non-empty"): + coco.prev_type_id(HostileString(" "), "Entry") + + module = "old:package.models" + qualname = "Outer.Source.Entry" + marker = cast(Any, coco.prev_type_id(module, qualname)) + + for attribute in ("module", "qualname"): + with pytest.raises(AttributeError): + setattr(marker, attribute, "mutated") + + variants: list[Any] = [ + marker, + copy.copy(marker), + copy.deepcopy(marker), + *( + pickle.loads(pickle.dumps(marker, protocol=protocol)) + for protocol in range(pickle.HIGHEST_PROTOCOL + 1) + ), + ] + for variant in variants: + assert type(variant) is type(marker) + assert (variant.module, variant.qualname) == (module, qualname) + + class MovedSourceEntry: + __coco_memo_type_id__: ClassVar[str] = variant + + assert _memo_fingerprint._type_identity_parts(MovedSourceEntry) == ( + module, + qualname, + ) + + class RegisteredSourceEntry: + pass + + try: + register_memo_key_function(RegisteredSourceEntry, stable_type_id=variants[-1]) + assert _memo_fingerprint._type_identity_parts(RegisteredSourceEntry) == ( + module, + qualname, + ) + finally: + unregister_memo_key_function(RegisteredSourceEntry) + + class OrdinaryStringSourceEntry: + __coco_memo_type_id__: ClassVar[str] = str(marker) + + assert _memo_fingerprint._type_identity_parts(OrdinaryStringSourceEntry) == ( + ("__coco_memo_type_id__", str(marker)), + None, + ) + + def test_pydantic_stable_type_id_allows_renamed_model_reuse() -> None: try: from pydantic import BaseModel @@ -472,6 +679,45 @@ def metaclass_state(cls: type, prev_state: object) -> MemoStateOutcome: unregister_memo_key_function(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_key_function(OldMemoMeta, metaclass_key) + register_memo_key_function( + 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_key_function(OldMemoMeta) + unregister_memo_key_function(MovedMemoMeta) + + def test_raw_class_object_honors_registered_type_memo_key_and_state() -> None: key_calls: list[type] = [] state_calls: list[tuple[type, object]] = [] @@ -622,6 +868,50 @@ class ChildB(Base): unregister_memo_key_function(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_key_function(OldBase, base_key) + register_memo_key_function( + 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_key_function(OldBase) + unregister_memo_key_function(MovedBase) + + def test_register_memo_key_function_registers_stable_type_id_without_key_function() -> ( None ): @@ -1199,6 +1489,18 @@ class Entry: pass assert coco.register_memo_key_function is register_memo_key_function + assert coco.prev_type_id is _memo_fingerprint.prev_type_id + assert "prev_type_id" in coco.__all__ + previous_type_id = coco.prev_type_id("old_package.models", "SourceEntry") + assert isinstance(previous_type_id, str) + with pytest.raises(TypeError, match="prev_type_id.*module must be a str"): + coco.prev_type_id(cast(Any, object()), "SourceEntry") + with pytest.raises(TypeError, match="prev_type_id.*qualname must be a str"): + coco.prev_type_id("old_package.models", cast(Any, object())) + with pytest.raises(ValueError, match="prev_type_id.*module must be non-empty"): + coco.prev_type_id("", "SourceEntry") + with pytest.raises(ValueError, match="prev_type_id.*qualname must be non-empty"): + coco.prev_type_id("old_package.models", " ") assert "register_memo_type_identifier" not in coco.__all__ assert not hasattr(coco, "register_memo_type_identifier") assert "register_not_memo_keyable" not in coco.__all__ From 9effa2374dac953f923519bb0aa8ff4abe87eeb5 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 18 Jul 2026 14:22:21 -0700 Subject: [PATCH 08/16] fix(memo): address type registration review feedback --- .../docs/advanced_topics/memoization_keys.mdx | 14 +- .../cocoindex/_internal/memo_fingerprint.py | 187 +++++++----------- .../tests/internal/test_memo_fingerprint.py | 52 ++++- 3 files changed, 129 insertions(+), 124 deletions(-) diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index a6c223d69..127790c96 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -5,7 +5,15 @@ description: > 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. @@ -24,8 +32,6 @@ Most values need no customization. CocoIndex already handles primitives, contain ## How data fingerprinting works -Timing depends on where the value comes from. Function arguments are fingerprinted when the function is called. Change-detected context values are fingerprinted when they are provided (for example, by `builder.provide()`), before `use_context()` reads them. `deps` values are fingerprinted once when the `@coco.fn(..., deps=...)` or `@coco.fn.as_async(..., deps=...)` decorator is applied. - 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. This applies to class instances, not class objects (`type`). @@ -161,7 +167,7 @@ Use `stable_type_id` without a key function only when CocoIndex already fingerpr Rule of thumb: if CocoIndex would otherwise use pickle for your object, `stable_type_id` by itself is not enough; register a key function and stable type ID together. -Omit the key function argument entirely when registering only a stable type ID; do not pass `None` as the second argument. +Either omit the key function argument or pass `None` when registering only a stable type ID; both forms are equivalent. ```python from dataclasses import dataclass diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index 50adc583f..9437ed304 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -32,21 +32,16 @@ _StateFn = typing.Callable[[typing.Any, typing.Any], typing.Any] -class _KeyFnUnset: - def __repr__(self) -> str: - return "" - - -_KEY_FN_UNSET = _KeyFnUnset() - - -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[int, tuple[weakref.ReferenceType[type], _MemoFns]] = {} -_stable_type_ids: dict[int, tuple[weakref.ReferenceType[type], str]] = {} +_STABLE_TYPE_ID_MISSING = object() +_memo_type_registry: dict[ + int, tuple[weakref.ReferenceType[type], _MemoTypeRegistry] +] = {} class StateFnEntry(typing.NamedTuple): @@ -206,98 +201,62 @@ def prev_type_id(module: str, qualname: str) -> str: return _PreviousTypeId(module, qualname) -def _remove_stable_type_id(type_id: int, dead_ref: weakref.ReferenceType[type]) -> None: +def _remove_memo_type_registry( + type_id: int, dead_ref: weakref.ReferenceType[type] +) -> None: """Remove ``type_id`` only if ``dead_ref`` is still the stored weakref.""" - entry = _stable_type_ids.get(type_id) + entry = _memo_type_registry.get(type_id) if entry is not None and entry[0] is dead_ref: - _stable_type_ids.pop(type_id, None) + _memo_type_registry.pop(type_id, None) -def _register_stable_type_id(typ: type, stable_type_id: str) -> None: - """Register a stable type ID for one exact Python type.""" +def _register_memo_type_registry(typ: type, registry: _MemoTypeRegistry) -> None: + """Register memo configuration for one exact Python type.""" type_id = id(typ) - def _remove_stale_type_id( - dead_ref: weakref.ReferenceType[type], - ) -> None: - _remove_stable_type_id(type_id, dead_ref) + def _remove_stale_registry(dead_ref: weakref.ReferenceType[type]) -> None: + _remove_memo_type_registry(type_id, dead_ref) - _stable_type_ids[type_id] = ( - weakref.ref(typ, _remove_stale_type_id), - stable_type_id, + _memo_type_registry[type_id] = ( + weakref.ref(typ, _remove_stale_registry), + registry, ) -def _unregister_stable_type_id(typ: type) -> None: - """Best-effort removal of an exact-type stable type ID registration.""" +def _unregister_memo_type_registry(typ: type) -> None: + """Best-effort removal of an exact-type memo registration.""" type_id = id(typ) - entry = _stable_type_ids.get(type_id) + entry = _memo_type_registry.get(type_id) if entry is not None and entry[0]() is typ: - _stable_type_ids.pop(type_id, None) + _memo_type_registry.pop(type_id, None) -def _registered_stable_type_id(typ: type) -> str | None: - """Return the registered stable type ID for ``typ`` from the id-keyed table.""" +def _registered_memo_type_registry(typ: type) -> _MemoTypeRegistry | None: + """Return the exact type's registration from the identity-keyed table.""" type_id = id(typ) - entry = _stable_type_ids.get(type_id) + entry = _memo_type_registry.get(type_id) if entry is None: return None - ref, stable_type_id = entry + ref, registry = entry if ref() is typ: - return stable_type_id - _stable_type_ids.pop(type_id, None) + return registry + _memo_type_registry.pop(type_id, None) return None def _lookup_stable_type_id(typ: type) -> str | None: """Resolve a registered or exact ``__coco_memo_type_id__`` stable type ID.""" - stable_type_id = _registered_stable_type_id(typ) - if stable_type_id is not None: - return stable_type_id - if "__coco_memo_type_id__" in typ.__dict__: - return _validate_stable_type_id( - typ.__dict__["__coco_memo_type_id__"], - source=f"{_memo_type_label(typ)}.__coco_memo_type_id__", - ) - return None - - -def _remove_memo_fns(type_id: int, dead_ref: weakref.ReferenceType[type]) -> None: - """Remove ``type_id`` only if ``dead_ref`` is still the stored weakref.""" - entry = _memo_fns.get(type_id) - if entry is not None and entry[0] is dead_ref: - _memo_fns.pop(type_id, None) - - -def _register_memo_fns(typ: type, memo_fns: _MemoFns) -> None: - """Register memo functions for one exact Python type.""" - type_id = id(typ) - - def _remove_stale_memo_fns(dead_ref: weakref.ReferenceType[type]) -> None: - _remove_memo_fns(type_id, dead_ref) + registry = _registered_memo_type_registry(typ) + if registry is not None and registry.stable_type_id is not None: + return registry.stable_type_id - _memo_fns[type_id] = (weakref.ref(typ, _remove_stale_memo_fns), memo_fns) - - -def _unregister_memo_fns(typ: type) -> None: - """Best-effort removal of an exact-type memo function registration.""" - type_id = id(typ) - entry = _memo_fns.get(type_id) - if entry is not None and entry[0]() is typ: - _memo_fns.pop(type_id, None) - - -def _registered_memo_fns(typ: type) -> _MemoFns | None: - """Return registered memo functions for ``typ`` from the id-keyed table.""" - type_id = id(typ) - entry = _memo_fns.get(type_id) - if entry is None: + stable_type_id = typ.__dict__.get("__coco_memo_type_id__", _STABLE_TYPE_ID_MISSING) + if stable_type_id is _STABLE_TYPE_ID_MISSING: return None - ref, memo_fns = entry - if ref() is typ: - return memo_fns - _memo_fns.pop(type_id, None) - return None + return _validate_stable_type_id( + stable_type_id, + source=f"{_memo_type_label(typ)}.__coco_memo_type_id__", + ) _MEMO_KEY_ATTR = "__coco_memo_key__" @@ -342,16 +301,18 @@ def _canonicalize_key_fragment( def _canonicalize_registered_memo_key( obj: object, owner: type, - memo: _MemoFns, + registry: _MemoTypeRegistry, state: _CanonicalizeState, state_methods: list[StateFnEntry], ) -> Fingerprintable: - key = memo.key_fn(obj) + key_fn = registry.key_fn + assert key_fn is not None + key = key_fn(obj) tag = "hook" - if memo.state_fn is not None: + if registry.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)) + bound = functools.partial(registry.state_fn, obj) + state_methods.append(_make_state_fn_entry(bound, registry.state_fn)) return ( tag, *_type_identity_parts(owner), @@ -370,10 +331,10 @@ def _canonicalize_class_object( for owner in metaclass.__mro__: if owner is object: break - memo = _registered_memo_fns(owner) - if memo is not None: + registry = _registered_memo_type_registry(owner) + if registry is not None and registry.key_fn is not None: return _canonicalize_registered_memo_key( - cls, owner, memo, state, state_methods + cls, owner, registry, state, state_methods ) return ( @@ -471,6 +432,7 @@ def register_memo_key_function( @typing.overload def register_memo_key_function( typ: type, + key_fn: None = None, *, stable_type_id: str, ) -> None: ... @@ -478,7 +440,7 @@ def register_memo_key_function( def register_memo_key_function( typ: type, - key_fn: _KeyFn | object = _KEY_FN_UNSET, + key_fn: _KeyFn | None = None, *, state_fn: _StateFn | None = None, stable_type_id: str | None = None, @@ -489,14 +451,14 @@ def register_memo_key_function( type wins. Stable type IDs registered without a key function apply to the exact type only; stable type IDs registered with a key function identify that selected owner type. Each call replaces the full registration for - ``typ``: omitting ``stable_type_id`` clears any previous stable type ID registered for ``typ``, - and omitting ``key_fn`` clears any previous key/state functions. - - To register only a stable type ID, omit ``key_fn`` rather than passing - ``None``. 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. + ``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. """ if not isinstance(typ, type): @@ -510,11 +472,6 @@ def register_memo_key_function( source="register_memo_key_function(..., stable_type_id)", ) if key_fn is None: - raise TypeError( - "register_memo_key_function() key_fn must be callable; omit key_fn " - "when registering only a stable type ID" - ) - if key_fn is _KEY_FN_UNSET: if state_fn is not None: raise TypeError( "register_memo_key_function() state_fn requires a memo key function" @@ -534,14 +491,14 @@ def register_memo_key_function( f"got {type(state_fn).__name__}" ) - if stable_type_id is not None: - _register_stable_type_id(typ, stable_type_id) - else: - _unregister_stable_type_id(typ) - if key_fn is not _KEY_FN_UNSET: - _register_memo_fns(typ, _MemoFns(typing.cast(_KeyFn, key_fn), state_fn)) - else: - _unregister_memo_fns(typ) + _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: @@ -566,16 +523,14 @@ def _raise_not_memo_keyable(obj: object) -> typing.NoReturn: "This type maintains internal state that is incompatible with memoization." ) - _unregister_stable_type_id(typ) - _register_memo_fns(typ, _MemoFns(_raise_not_memo_keyable)) + _register_memo_type_registry(typ, _MemoTypeRegistry(key_fn=_raise_not_memo_keyable)) def unregister_memo_key_function(typ: type) -> None: """Remove registered memo key function and stable type ID (best-effort).""" if isinstance(typ, type): - _unregister_stable_type_id(typ) - _unregister_memo_fns(typ) + _unregister_memo_type_registry(typ) def _stable_sort_key(v: Fingerprintable) -> tuple[typing.Any, ...]: @@ -648,10 +603,10 @@ def _canonicalize( ) for owner in type(obj).__mro__: - memo = _registered_memo_fns(owner) - if memo is not None: + registry = _registered_memo_type_registry(owner) + if registry is not None and registry.key_fn is not None: return _canonicalize_registered_memo_key( - obj, owner, memo, state, state_methods + obj, owner, registry, state, state_methods ) # 3) Cycle / shared-reference tracking diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index b3a56943c..6cd5cbf82 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -983,6 +983,30 @@ def __coco_memo_key__(self) -> object: unregister_memo_key_function(SameStableTypeId) +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_key_function(Parent, lambda entry: ("parent", entry.value)) + register_memo_key_function(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_key_function(Child) + unregister_memo_key_function(Parent) + + def test_stable_type_id_only_registration_replaces_key_and_state_functions() -> None: @dataclasses.dataclass class Entry: @@ -1226,14 +1250,26 @@ class SameStableTypeId: unregister_memo_key_function(SameStableTypeId) -def test_register_memo_key_function_rejects_explicit_none_key_function() -> None: +def test_register_memo_key_function_accepts_explicit_none_key_function() -> None: class Entry: - pass + def __coco_memo_key__(self) -> object: + return ("entry",) + + class SameStableTypeId: + def __coco_memo_key__(self) -> object: + return ("entry",) - with pytest.raises(TypeError, match="key_fn"): + try: + register_memo_key_function(Entry, None, stable_type_id="test.ExplicitNone/v1") register_memo_key_function( - Entry, cast(Any, None), stable_type_id="test.ExplicitNone/v1" + SameStableTypeId, stable_type_id="test.ExplicitNone/v1" + ) + assert fingerprint_call(_dummy_fn, (Entry(),), {}, []) == fingerprint_call( + _dummy_fn, (SameStableTypeId(),), {}, [] ) + finally: + unregister_memo_key_function(Entry) + unregister_memo_key_function(SameStableTypeId) @pytest.mark.parametrize( @@ -1467,6 +1503,12 @@ class EmptyId: def __coco_memo_key__(self) -> object: return ("bad", 1) + class NoneId: + __coco_memo_type_id__ = None + + def __coco_memo_key__(self) -> object: + return ("bad", 1) + assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( _dummy_fn, (NewEntry(1),), {}, [] ) @@ -1478,6 +1520,8 @@ def __coco_memo_key__(self) -> object: ) with pytest.raises(TypeError, match="must be a str"): fingerprint_call(_dummy_fn, (BadObjectId(),), {}, []) + with pytest.raises(TypeError, match="must be a str"): + fingerprint_call(_dummy_fn, (NoneId(),), {}, []) with pytest.raises(ValueError, match="non-empty"): fingerprint_call(_dummy_fn, (EmptyId(),), {}, []) From d0f088081f24df4b221a2345c9970584f71edc9e Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 25 Jul 2026 10:35:44 -0700 Subject: [PATCH 09/16] fix(memo): address final type identity review feedback --- .../docs/advanced_topics/memoization_keys.mdx | 4 +- .../cocoindex/_internal/memo_fingerprint.py | 102 ++++--- .../tests/internal/test_memo_fingerprint.py | 262 ++++++++++++++---- 3 files changed, 257 insertions(+), 111 deletions(-) diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index 127790c96..696def2b1 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -242,7 +242,7 @@ Registration is replacement-based, not additive. Each call replaces the full reg | `register_memo_key_function(T, key_fn)` | `key_fn` is installed and any previous `state_fn` is cleared unless a new `state_fn` is supplied. | Any previous registered stable type ID is cleared. | | `register_memo_key_function(T, stable_type_id="...")` | Any previous key/state functions are cleared. | The provided stable type ID is installed. | -A `__coco_memo_type_id__` defined on the class is separate from the process-global registration table. If both are set for the same exact type, the registered stable type ID takes precedence. Provide `stable_type_id=` in a registration only when you intentionally want that override; otherwise, register only the key/state function and use the class-defined ID. +A `__coco_memo_type_id__` defined on the class is separate from the process-global registration table. If both are set for the same exact type, the class-defined stable type ID takes precedence. The registration's `stable_type_id=` is used only when the selected owner has no non-`None` class-defined ID. For a registry-owned stable type ID, if a type needs a key function, state function, and stable type ID, register all of them in the same call. Avoid splitting registration across helper modules where import order can determine which registration wins. @@ -277,7 +277,7 @@ coco.register_memo_key_function(Parent, stable_type_id="com.example.Parent/v1") #### Stable type ID with a key function: owner selected by MRO -When `register_memo_key_function(Parent, key_fn, stable_type_id=...)` handles a `Child(Parent)` value through MRO lookup, the selected registered owner (`Parent`) supplies the type namespace. That owner namespace can come from the registration's `stable_type_id=...` argument or from `Parent.__coco_memo_type_id__`. For example: +When `register_memo_key_function(Parent, key_fn, stable_type_id=...)` handles a `Child(Parent)` value through MRO lookup, the selected registered owner (`Parent`) supplies the type namespace. `Parent.__coco_memo_type_id__` supplies that namespace when it is non-`None`; otherwise, the registration's `stable_type_id=...` is used. For example: ```python class Parent: diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index 9437ed304..928e1125b 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -38,7 +38,6 @@ class _MemoTypeRegistry(typing.NamedTuple): stable_type_id: str | None = None -_STABLE_TYPE_ID_MISSING = object() _memo_type_registry: dict[ int, tuple[weakref.ReferenceType[type], _MemoTypeRegistry] ] = {} @@ -149,29 +148,19 @@ def _memo_type_label(typ: type) -> str: class _PreviousTypeId(str): """A prior automatic type identity carried through the stable-ID path.""" - __slots__ = () - - def __new__(cls, module: str, qualname: str | None = None) -> _PreviousTypeId: - payload = module if qualname is None else f"{len(module)}:{module}{qualname}" - return super().__new__(cls, payload) - - def __getnewargs__(self) -> tuple[str]: - return (str(self),) + __slots__ = ("_identity_parts",) + _identity_parts: tuple[str, str] - def _identity_parts(self) -> tuple[str, str]: - module_length_str, separator, payload = self.partition(":") - if separator == "": - raise ValueError("invalid previous type identity payload") - module_length = int(module_length_str) - return payload[:module_length], payload[module_length:] + def __new__(cls, module: str, qualname: str) -> _PreviousTypeId: + marker = super().__new__(cls, f"{len(module)}:{module}{qualname}") + object.__setattr__(marker, "_identity_parts", (module, qualname)) + return marker - @property - def module(self) -> str: - return self._identity_parts()[0] + def __setattr__(self, name: str, value: object) -> typing.NoReturn: + raise AttributeError(f"{type(self).__name__} is immutable") - @property - def qualname(self) -> str: - return self._identity_parts()[1] + def __reduce__(self) -> tuple[type[_PreviousTypeId], tuple[str, str]]: + return type(self), self._identity_parts def _validate_stable_type_id(stable_type_id: object, *, source: str) -> str: @@ -244,21 +233,6 @@ def _registered_memo_type_registry(typ: type) -> _MemoTypeRegistry | None: return None -def _lookup_stable_type_id(typ: type) -> str | None: - """Resolve a registered or exact ``__coco_memo_type_id__`` stable type ID.""" - registry = _registered_memo_type_registry(typ) - if registry is not None and registry.stable_type_id is not None: - return registry.stable_type_id - - stable_type_id = typ.__dict__.get("__coco_memo_type_id__", _STABLE_TYPE_ID_MISSING) - if stable_type_id is _STABLE_TYPE_ID_MISSING: - return None - return _validate_stable_type_id( - stable_type_id, - source=f"{_memo_type_label(typ)}.__coco_memo_type_id__", - ) - - _MEMO_KEY_ATTR = "__coco_memo_key__" _MEMO_STATE_ATTR = "__coco_memo_state__" _CLASS_OBJECT_OWNER_IDENTITY: tuple[Fingerprintable, Fingerprintable] = ( @@ -267,7 +241,10 @@ def _lookup_stable_type_id(typ: type) -> str | None: ) -def _type_identity_parts(typ: type) -> tuple[Fingerprintable, Fingerprintable]: +def _type_identity_parts( + typ: type, + registry: _MemoTypeRegistry | None, +) -> 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 @@ -275,9 +252,17 @@ def _type_identity_parts(typ: type) -> tuple[Fingerprintable, Fingerprintable]: tagged first slot keeps stable type IDs disjoint from ordinary module names; ``None`` fills the qualname slot. """ - stable_type_id = _lookup_stable_type_id(typ) + stable_type_id = typ.__dict__.get("__coco_memo_type_id__") + if stable_type_id is not None: + stable_type_id = _validate_stable_type_id( + stable_type_id, + source=f"{_memo_type_label(typ)}.__coco_memo_type_id__", + ) + elif registry is not None: + stable_type_id = registry.stable_type_id + if isinstance(stable_type_id, _PreviousTypeId): - return stable_type_id._identity_parts() + return stable_type_id._identity_parts if stable_type_id is not None: return (("__coco_memo_type_id__", stable_type_id), None) return (canonical_module_name(typ), getattr(typ, "__qualname__", None)) @@ -315,7 +300,7 @@ def _canonicalize_registered_memo_key( state_methods.append(_make_state_fn_entry(bound, registry.state_fn)) return ( tag, - *_type_identity_parts(owner), + *_type_identity_parts(owner, registry), _canonicalize_key_fragment(key, state, state_methods), ) @@ -328,21 +313,29 @@ def _canonicalize_class_object( """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 = _registered_memo_type_registry(owner) + 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)), + ("seq", _type_identity_parts(cls, cls_registry)), ) @@ -358,6 +351,7 @@ def _is_pydantic_model(obj: object) -> bool: def _canonicalize_dataclass( obj: object, + registry: _MemoTypeRegistry | None, state: _CanonicalizeState, state_methods: list[StateFnEntry], ) -> Fingerprintable: @@ -370,7 +364,7 @@ def _canonicalize_dataclass( fields = dataclasses.fields(obj) # type: ignore[arg-type] return ( "dataclass", - *_type_identity_parts(typ), + *_type_identity_parts(typ, registry), tuple( (field.name, _canonicalize(getattr(obj, field.name), state, state_methods)) for field in fields @@ -380,6 +374,7 @@ def _canonicalize_dataclass( def _canonicalize_pydantic( obj: object, + registry: _MemoTypeRegistry | None, state: _CanonicalizeState, state_methods: list[StateFnEntry], ) -> Fingerprintable: @@ -392,7 +387,7 @@ def _canonicalize_pydantic( field_names = obj.__pydantic_fields__.keys() # type: ignore[attr-defined] return ( "pydantic", - *_type_identity_parts(typ), + *_type_identity_parts(typ, registry), tuple( (name, _canonicalize(getattr(obj, name), state, state_methods)) for name in field_names @@ -585,10 +580,11 @@ def _canonicalize( 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, _MEMO_STATE_ATTR, None) if state_hook is not None and callable(state_hook): @@ -598,15 +594,17 @@ def _canonicalize( state_methods.append(_make_state_fn_entry(state_hook, raw_fn)) return ( tag, - *_type_identity_parts(typ), + *_type_identity_parts(typ, registry), _canonicalize_key_fragment(k, state, state_methods), ) - for owner in type(obj).__mro__: - registry = _registered_memo_type_registry(owner) - if registry is not None and registry.key_fn is not None: + 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, registry, state, state_methods + obj, owner, owner_registry, state, state_methods ) # 3) Cycle / shared-reference tracking @@ -640,11 +638,11 @@ def _canonicalize( # 5) Dataclass instances if _is_dataclass_instance(obj): - return _canonicalize_dataclass(obj, state, state_methods) + return _canonicalize_dataclass(obj, registry, state, state_methods) # 6) Pydantic v2 models if _is_pydantic_model(obj): - return _canonicalize_pydantic(obj, state, state_methods) + return _canonicalize_pydantic(obj, registry, state, state_methods) # 7) Fallback try: diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index 6cd5cbf82..2eb84ffe3 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -313,13 +313,8 @@ def __format__(self, format_spec: str) -> str: def strip(self, chars: str | None = None) -> str: raise AssertionError("subclass strip should not be called") - normalized_marker = cast( - Any, - coco.prev_type_id(HostileString("old.package"), HostileString("Outer.Entry")), - ) - assert (normalized_marker.module, normalized_marker.qualname) == ( - "old.package", - "Outer.Entry", + normalized_marker = coco.prev_type_id( + HostileString("old.package"), HostileString("Outer.Entry") ) assert normalized_marker == coco.prev_type_id("old.package", "Outer.Entry") with pytest.raises(ValueError, match="non-empty"): @@ -327,13 +322,21 @@ def strip(self, chars: str | None = None) -> str: module = "old:package.models" qualname = "Outer.Source.Entry" - marker = cast(Any, coco.prev_type_id(module, qualname)) + marker = coco.prev_type_id(module, qualname) + assert isinstance(marker, _memo_fingerprint._PreviousTypeId) + + class IdentitySourceEntry: + __coco_memo_type_id__: ClassVar[str] = marker - for attribute in ("module", "qualname"): - with pytest.raises(AttributeError): + assert ( + _memo_fingerprint._type_identity_parts(IdentitySourceEntry, None) + is marker._identity_parts + ) + for attribute in ("_identity_parts", "extra"): + with pytest.raises(AttributeError, match="immutable"): setattr(marker, attribute, "mutated") - variants: list[Any] = [ + variants = [ marker, copy.copy(marker), copy.deepcopy(marker), @@ -344,12 +347,11 @@ def strip(self, chars: str | None = None) -> str: ] for variant in variants: assert type(variant) is type(marker) - assert (variant.module, variant.qualname) == (module, qualname) class MovedSourceEntry: __coco_memo_type_id__: ClassVar[str] = variant - assert _memo_fingerprint._type_identity_parts(MovedSourceEntry) == ( + assert _memo_fingerprint._type_identity_parts(MovedSourceEntry, None) == ( module, qualname, ) @@ -359,18 +361,23 @@ class RegisteredSourceEntry: try: register_memo_key_function(RegisteredSourceEntry, stable_type_id=variants[-1]) - assert _memo_fingerprint._type_identity_parts(RegisteredSourceEntry) == ( - module, - qualname, + 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_key_function(RegisteredSourceEntry) + ordinary_marker = str(marker) + class OrdinaryStringSourceEntry: - __coco_memo_type_id__: ClassVar[str] = str(marker) + __coco_memo_type_id__: ClassVar[str] = ordinary_marker - assert _memo_fingerprint._type_identity_parts(OrdinaryStringSourceEntry) == ( - ("__coco_memo_type_id__", str(marker)), + assert _memo_fingerprint._type_identity_parts(OrdinaryStringSourceEntry, None) == ( + ("__coco_memo_type_id__", ordinary_marker), None, ) @@ -485,7 +492,7 @@ def make_graph() -> list[object]: assert canonical == ( "seq", - (("hook", *_memo_fingerprint._type_identity_parts(Entry), ("ref", 0)),), + (("hook", *_memo_fingerprint._type_identity_parts(Entry, None), ("ref", 0)),), ) assert _memo_fingerprint.memo_fingerprint( graph @@ -509,7 +516,13 @@ def make_graph() -> list[object]: assert canonical == ( "seq", - (("hook", *_memo_fingerprint._type_identity_parts(Entry), ("ref", 0)),), + ( + ( + "hook", + *_memo_fingerprint._type_identity_parts(Entry, None), + ("ref", 0), + ), + ), ) assert _memo_fingerprint.memo_fingerprint( graph @@ -544,12 +557,12 @@ def __coco_memo_key__(self) -> object: ( ( "hook", - *_memo_fingerprint._type_identity_parts(FirstEntry), + *_memo_fingerprint._type_identity_parts(FirstEntry, None), ("seq", ("first",)), ), ( "hook", - *_memo_fingerprint._type_identity_parts(SecondEntry), + *_memo_fingerprint._type_identity_parts(SecondEntry, None), ("seq", ("second",)), ), ), @@ -569,7 +582,7 @@ def __coco_memo_key__(self) -> object: assert top_level == ( "hook", - *_memo_fingerprint._type_identity_parts(Entry), + *_memo_fingerprint._type_identity_parts(Entry, None), ("seq", (("seq", ("shared",)), ("ref", 1))), ) assert parent_wrapped == ( @@ -577,7 +590,7 @@ def __coco_memo_key__(self) -> object: ( ( "hook", - *_memo_fingerprint._type_identity_parts(Entry), + *_memo_fingerprint._type_identity_parts(Entry, None), ("seq", (("seq", ("shared",)), ("ref", 2))), ), ), @@ -1067,56 +1080,173 @@ def __init__(self, value: object) -> None: unregister_memo_key_function(SameStableTypeId) -def test_key_only_registration_falls_back_to_class_declared_stable_type_id() -> None: - class Entry: - __coco_memo_type_id__ = "test.DeclaredFallback/v1" +def test_intrinsic_hooks_and_declared_stable_id_beat_registration() -> None: + declared_type_id = "test.DeclaredIntrinsic/v1" - def __init__(self, value: object) -> None: - self.value = value + class Entry: + __coco_memo_type_id__ = declared_type_id def __coco_memo_key__(self) -> object: - return ("entry", self.value) + return "intrinsic-key" - class DeclaredPeer: - __coco_memo_type_id__ = "test.DeclaredFallback/v1" + def __coco_memo_state__(self, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome(state=prev_state, memo_valid=True) - def __init__(self, value: object) -> None: - self.value = value + def registered_key(_entry: Entry) -> object: + raise AssertionError("intrinsic key must beat the registered key") - def __coco_memo_key__(self) -> object: - return ("entry", self.value) + def registered_state(_entry: Entry, _prev_state: object) -> MemoStateOutcome: + raise AssertionError("intrinsic state must beat the registered state") + + try: + register_memo_key_function( + 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_key_function(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 - class RegisteredPeer: def __init__(self, value: object) -> None: self.value = value - def __coco_memo_key__(self) -> object: - return ("entry", self.value) - - assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) == fingerprint_call( - _dummy_fn, (DeclaredPeer(1),), {}, [] - ) + class ChildEntry(BaseEntry): + pass try: - register_memo_key_function(Entry, stable_type_id="test.RegisteredOverride/v1") register_memo_key_function( - RegisteredPeer, stable_type_id="test.RegisteredOverride/v1" + 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_key_function(BaseEntry) + + +def test_identity_dispatch_looks_up_each_candidate_registry_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class ExactEntry: + pass + + class RegisteredBase: + pass + + class InheritedEntry(RegisteredBase): + pass + + @dataclasses.dataclass + class DataclassEntry: + value: int + + class MemoMeta(type): + pass - assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) == fingerprint_call( - _dummy_fn, (RegisteredPeer(1),), {}, [] + class MetaEntry(metaclass=MemoMeta): + pass + + original_lookup = _memo_fingerprint._registered_memo_type_registry + lookups: list[type] = [] + + def tracked_lookup(typ: type) -> Any: + lookups.append(typ) + return original_lookup(typ) + + def assert_lookups(*expected: type) -> None: + assert lookups == list(expected) + lookups.clear() + + try: + register_memo_key_function(ExactEntry, lambda _entry: "exact") + register_memo_key_function(RegisteredBase, lambda _entry: "inherited") + register_memo_key_function( + DataclassEntry, stable_type_id="test.LookupDataclass/v1" ) - assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) != fingerprint_call( - _dummy_fn, (DeclaredPeer(1),), {}, [] + register_memo_key_function(MemoMeta, lambda _cls: "metaclass") + monkeypatch.setattr( + _memo_fingerprint, "_registered_memo_type_registry", tracked_lookup ) - register_memo_key_function(Entry, lambda entry: ("entry", entry.value)) - assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) == fingerprint_call( - _dummy_fn, (DeclaredPeer(1),), {}, [] + _memo_fingerprint._canonicalize(ExactEntry(), None, []) + assert_lookups(ExactEntry) + + _memo_fingerprint._canonicalize(InheritedEntry(), None, []) + assert_lookups(InheritedEntry, RegisteredBase) + + _memo_fingerprint._canonicalize(DataclassEntry(1), None, []) + assert_lookups(DataclassEntry, object) + + _memo_fingerprint._canonicalize(MetaEntry, None, []) + assert_lookups(MemoMeta) + finally: + unregister_memo_key_function(ExactEntry) + unregister_memo_key_function(RegisteredBase) + unregister_memo_key_function(DataclassEntry) + unregister_memo_key_function(MemoMeta) + + +def test_pydantic_identity_registry_lookup_is_not_repeated( + monkeypatch: pytest.MonkeyPatch, +) -> None: + try: + from pydantic import BaseModel + except ImportError: + pytest.skip("pydantic not installed") + return + + class Model(BaseModel): + value: int + + original_lookup = _memo_fingerprint._registered_memo_type_registry + lookups: list[type] = [] + + def tracked_lookup(typ: type) -> Any: + lookups.append(typ) + return original_lookup(typ) + + try: + register_memo_key_function(Model, stable_type_id="test.LookupPydantic/v1") + monkeypatch.setattr( + _memo_fingerprint, "_registered_memo_type_registry", tracked_lookup + ) + canonical = _memo_fingerprint._canonicalize(Model(value=1), None, []) + assert isinstance(canonical, tuple) + + assert canonical[:3] == ( + "pydantic", + ("__coco_memo_type_id__", "test.LookupPydantic/v1"), + None, ) + assert sum(owner is Model for owner in lookups) == 1 + assert len({id(owner) for owner in lookups}) == len(lookups) finally: - unregister_memo_key_function(Entry) - unregister_memo_key_function(RegisteredPeer) + unregister_memo_key_function(Model) def test_combined_registration_uses_stable_type_id_and_collects_state_fn() -> None: @@ -1520,8 +1650,26 @@ def __coco_memo_key__(self) -> object: ) with pytest.raises(TypeError, match="must be a str"): fingerprint_call(_dummy_fn, (BadObjectId(),), {}, []) - with pytest.raises(TypeError, match="must be a str"): - fingerprint_call(_dummy_fn, (NoneId(),), {}, []) + 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_key_function(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_key_function(NoneId) with pytest.raises(ValueError, match="non-empty"): fingerprint_call(_dummy_fn, (EmptyId(),), {}, []) From 1744a018a593c5fcf279054df23876b9466d653d Mon Sep 17 00:00:00 2001 From: Max Rong Date: Fri, 28 Aug 2026 06:45:02 +0000 Subject: [PATCH 10/16] feat(memo): simplify previous type identity and remove defensive validation - Simplify prev_type_id constructor and remove defensive validation helpers. - Enforce _PreviousTypeId immutability with __delattr__ alongside __setattr__. - Remove runtime stable type ID validation from canonicalization hot path. - Remove redundant _canonicalize_key_fragment wrapper. - Clean up and streamline _PreviousTypeId immutability and round-trip tests. --- .../cocoindex/_internal/memo_fingerprint.py | 72 +++------------- .../tests/internal/test_memo_fingerprint.py | 85 ++----------------- 2 files changed, 20 insertions(+), 137 deletions(-) diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index 928e1125b..ce5eeb42e 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -27,7 +27,6 @@ ) from .typing import Fingerprintable - _KeyFn = typing.Callable[[typing.Any], typing.Any] _StateFn = typing.Callable[[typing.Any, typing.Any], typing.Any] @@ -93,7 +92,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) @@ -140,18 +139,13 @@ def canonical_module_name(obj: typing.Any) -> str: return mod -def _memo_type_label(typ: type) -> str: - """Return a user-facing label for stable type ID validation errors.""" - return f"{canonical_module_name(typ)}.{getattr(typ, '__qualname__', '')}" - - 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) -> _PreviousTypeId: + 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 @@ -159,34 +153,15 @@ def __new__(cls, module: str, qualname: str) -> _PreviousTypeId: 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 _validate_stable_type_id(stable_type_id: object, *, source: str) -> str: - """Validate a non-empty stable type ID.""" - if not isinstance(stable_type_id, str): - raise TypeError(f"{source} must be a str, got {type(stable_type_id).__name__}") - if stable_type_id.strip() == "": - raise ValueError( - f"{source} must be non-empty and contain non-whitespace characters" - ) - return stable_type_id - - -def _validate_previous_type_id_part(value: object, *, source: str) -> str: - """Validate and normalize one previous automatic identity part.""" - if isinstance(value, str): - return _validate_stable_type_id(str.__str__(value), source=source) - return _validate_stable_type_id(value, source=source) - - def prev_type_id(module: str, qualname: str) -> str: """Return a marker that reuses a type's prior automatic identity.""" - module = _validate_previous_type_id_part(module, source="prev_type_id() module") - qualname = _validate_previous_type_id_part( - qualname, source="prev_type_id() qualname" - ) return _PreviousTypeId(module, qualname) @@ -253,12 +228,7 @@ def _type_identity_parts( names; ``None`` fills the qualname slot. """ stable_type_id = typ.__dict__.get("__coco_memo_type_id__") - if stable_type_id is not None: - stable_type_id = _validate_stable_type_id( - stable_type_id, - source=f"{_memo_type_label(typ)}.__coco_memo_type_id__", - ) - elif registry is not None: + if stable_type_id is None and registry is not None: stable_type_id = registry.stable_type_id if isinstance(stable_type_id, _PreviousTypeId): @@ -268,21 +238,6 @@ def _type_identity_parts( return (canonical_module_name(typ), getattr(typ, "__qualname__", None)) -def _canonicalize_key_fragment( - obj: object, - state: _CanonicalizeState, - state_methods: list[StateFnEntry], -) -> Fingerprintable: - """Canonicalize a memo-key fragment within the current root traversal. - - Sharing traversal state preserves cycles through the parent object and keeps - temporary fragment objects alive so their IDs cannot be reused during this - traversal. - """ - - return _canonicalize(obj, state, state_methods) - - def _canonicalize_registered_memo_key( obj: object, owner: type, @@ -301,7 +256,7 @@ def _canonicalize_registered_memo_key( return ( tag, *_type_identity_parts(owner, registry), - _canonicalize_key_fragment(key, state, state_methods), + _canonicalize(key, state, state_methods), ) @@ -461,11 +416,6 @@ def register_memo_key_function( "register_memo_key_function() expects typ to be a type, " f"got {type(typ).__name__}" ) - if stable_type_id is not None: - stable_type_id = _validate_stable_type_id( - stable_type_id, - source="register_memo_key_function(..., stable_type_id)", - ) if key_fn is None: if state_fn is not None: raise TypeError( @@ -595,7 +545,7 @@ def _canonicalize( return ( tag, *_type_identity_parts(typ, registry), - _canonicalize_key_fragment(k, state, state_methods), + _canonicalize(k, state, state_methods), ) for owner in typ.__mro__: @@ -649,7 +599,7 @@ def _canonicalize( 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." @@ -729,10 +679,10 @@ def fingerprint_call( __all__ = [ "NotMemoKeyable", + "fingerprint_call", + "memo_fingerprint", "prev_type_id", "register_memo_key_function", "register_not_memo_keyable", "unregister_memo_key_function", - "fingerprint_call", - "memo_fingerprint", ] diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index 2eb84ffe3..21f1e54c6 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -7,7 +7,6 @@ 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 ( @@ -27,14 +26,6 @@ def _dummy_fn(*args: Any, **kwargs: Any) -> None: raise RuntimeError("not called") -def _canonical_contains(value: object, needle: object) -> bool: - if value == needle: - return True - if isinstance(value, tuple): - return any(_canonical_contains(item, needle) for item in value) - return False - - def test_fingerprint_dict_order_independent() -> None: a = {"x": 1, "y": 2} b = {"y": 2, "x": 1} @@ -298,27 +289,6 @@ def test_prev_type_id_marker_is_immutable_and_round_trips() -> None: right = coco.prev_type_id("a", "b.C") assert str(left) != str(right) assert left != right - assert len({left, right}) == 2 - - class HostileString(str): - def __len__(self) -> int: - return 1 - - def __str__(self) -> str: - return "wrong-str" - - def __format__(self, format_spec: str) -> str: - return "wrong-format" - - def strip(self, chars: str | None = None) -> str: - raise AssertionError("subclass strip should not be called") - - normalized_marker = coco.prev_type_id( - HostileString("old.package"), HostileString("Outer.Entry") - ) - assert normalized_marker == coco.prev_type_id("old.package", "Outer.Entry") - with pytest.raises(ValueError, match="non-empty"): - coco.prev_type_id(HostileString(" "), "Entry") module = "old:package.models" qualname = "Outer.Source.Entry" @@ -332,21 +302,17 @@ class IdentitySourceEntry: _memo_fingerprint._type_identity_parts(IdentitySourceEntry, None) is marker._identity_parts ) - for attribute in ("_identity_parts", "extra"): - with pytest.raises(AttributeError, match="immutable"): - setattr(marker, attribute, "mutated") + with pytest.raises(AttributeError, match="immutable"): + marker._identity_parts = ("mutated", "mutated") + with pytest.raises(AttributeError, match="immutable"): + delattr(marker, "_identity_parts") - variants = [ + for variant in ( marker, copy.copy(marker), copy.deepcopy(marker), - *( - pickle.loads(pickle.dumps(marker, protocol=protocol)) - for protocol in range(pickle.HIGHEST_PROTOCOL + 1) - ), - ] - for variant in variants: - assert type(variant) is type(marker) + pickle.loads(pickle.dumps(marker)), + ): class MovedSourceEntry: __coco_memo_type_id__: ClassVar[str] = variant @@ -360,7 +326,7 @@ class RegisteredSourceEntry: pass try: - register_memo_key_function(RegisteredSourceEntry, stable_type_id=variants[-1]) + register_memo_key_function(RegisteredSourceEntry, stable_type_id=marker) registry = _memo_fingerprint._registered_memo_type_registry( RegisteredSourceEntry ) @@ -1593,7 +1559,7 @@ def __getstate__(self) -> object: unregister_memo_key_function(Registered) -def test_stable_type_id_exact_type_and_validation() -> None: +def test_stable_type_id_exact_type() -> None: class OldEntry: __coco_memo_type_id__ = "test.DirectEntry/v1" @@ -1621,18 +1587,6 @@ def __coco_memo_key__(self) -> object: class Child(Parent): pass - class BadObjectId: - __coco_memo_type_id__ = object() - - def __coco_memo_key__(self) -> object: - return ("bad", 1) - - class EmptyId: - __coco_memo_type_id__ = "" - - def __coco_memo_key__(self) -> object: - return ("bad", 1) - class NoneId: __coco_memo_type_id__ = None @@ -1648,8 +1602,6 @@ def __coco_memo_key__(self) -> object: assert fingerprint_call(_dummy_fn, (Parent(),), {}, []) != fingerprint_call( _dummy_fn, (Child(),), {}, [] ) - with pytest.raises(TypeError, match="must be a str"): - fingerprint_call(_dummy_fn, (BadObjectId(),), {}, []) none_id_canonical = _memo_fingerprint._canonicalize(NoneId(), None, []) assert isinstance(none_id_canonical, tuple) assert none_id_canonical[:3] == ( @@ -1670,29 +1622,16 @@ def __coco_memo_key__(self) -> object: ) finally: unregister_memo_key_function(NoneId) - with pytest.raises(ValueError, match="non-empty"): - fingerprint_call(_dummy_fn, (EmptyId(),), {}, []) def test_register_memo_key_function_validation_and_public_export() -> None: import cocoindex as coco - class Entry: - pass - assert coco.register_memo_key_function is register_memo_key_function assert coco.prev_type_id is _memo_fingerprint.prev_type_id assert "prev_type_id" in coco.__all__ previous_type_id = coco.prev_type_id("old_package.models", "SourceEntry") assert isinstance(previous_type_id, str) - with pytest.raises(TypeError, match="prev_type_id.*module must be a str"): - coco.prev_type_id(cast(Any, object()), "SourceEntry") - with pytest.raises(TypeError, match="prev_type_id.*qualname must be a str"): - coco.prev_type_id("old_package.models", cast(Any, object())) - with pytest.raises(ValueError, match="prev_type_id.*module must be non-empty"): - coco.prev_type_id("", "SourceEntry") - with pytest.raises(ValueError, match="prev_type_id.*qualname must be non-empty"): - coco.prev_type_id("old_package.models", " ") assert "register_memo_type_identifier" not in coco.__all__ assert not hasattr(coco, "register_memo_type_identifier") assert "register_not_memo_keyable" not in coco.__all__ @@ -1701,12 +1640,6 @@ class Entry: register_memo_key_function( cast(Any, object()), stable_type_id="test.Invalid/v1" ) - with pytest.raises(TypeError, match="must be a str"): - register_memo_key_function(Entry, stable_type_id=cast(Any, object())) - with pytest.raises(ValueError, match="non-empty"): - register_memo_key_function(Entry, stable_type_id="") - with pytest.raises(ValueError, match="non-empty"): - register_memo_key_function(Entry, stable_type_id=" ") def test_cycles_are_supported_and_deterministic() -> None: From 4ef25b92554bcefac48d9af4725618783f2ed991 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Fri, 28 Aug 2026 06:45:28 +0000 Subject: [PATCH 11/16] docs(memo): prune stable type ID guide and clarify class object precedence - Condense stable type IDs guide down to core migration patterns (proactive string IDs, prev_type_id, and third-party registration). - Clarify that registered custom metaclass or type key functions take precedence over __coco_memo_type_id__. - Clean up em-dash formatting and remove verbose MRO replacement tables. --- .../docs/advanced_topics/memoization_keys.mdx | 177 ++---------------- 1 file changed, 17 insertions(+), 160 deletions(-) diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index 696def2b1..df0aba22b 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -113,35 +113,24 @@ register_memo_key_function(Path, path_key) ### Use stable type IDs across refactors -This is an advanced migration feature. If you are not preserving existing memo entries across a class move or rename, skip this section and continue with [`memo_key=`](#override-at-the-call-site-with-memo_key) or [memo state validation](#memo-state-validation). - -By default, CocoIndex includes a type namespace (module plus qualified name) for dataclass instances, Pydantic model instances, objects handled by `__coco_memo_key__` or a registered memo key function, and class objects passed as values. If a class moves or is renamed, that default type namespace changes and old memo entries are not reused. - -A stable type ID replaces that type-name portion when the type should keep the same memo namespace 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 SourceEntry: - __coco_memo_type_id__: ClassVar[str] = "com.example.SourceEntry/v1" +class NewEntry: + __coco_memo_type_id__: ClassVar[str] = "com.example.NewEntry/v1" id: str content: str -``` - -**Choose a durable string ID proactively.** For a new type, or before any memo entries for it matter, use a globally unique durable string `__coco_memo_type_id__` as shown above when memo stability is important and a future move or rename cannot be ruled out. Introducing or changing an ordinary string ID after entries exist creates a new memo namespace and causes one cold execution for affected memoized calls. - -**Retain an existing automatic identity.** If a type already has reusable entries under its automatic identity, assign that exact previous identity to the moved class: - -```python -from dataclasses import dataclass -from typing import ClassVar -import cocoindex as coco +# Retain existing automatic entries: @dataclass -class SourceEntry: +class SourceEntry: # Moved to new_package.models __coco_memo_type_id__: ClassVar[str] = coco.prev_type_id( "old_package.models", "SourceEntry", @@ -151,77 +140,7 @@ class SourceEntry: version: int ``` -The arguments to `prev_type_id` are the exact previous canonical module and qualified name. Nested classes retain their enclosing names in the qualified name, such as `"Container.SourceEntry"`. A file-backed class previously defined in `__main__` uses the entry script's basename without `.py`; for example, `main.py` becomes `"main"`. - -`prev_type_id` does not scan, copy, promote, or migrate records. For a type-aware canonical form that already included the type's automatic identity, it keeps emitting the same `(old_module, old_qualname)`, so those existing entries remain directly addressable. Entries created through pickle fallback did not include this identity. Adding `prev_type_id` alone has no effect on pickle-fallback canonicalization and cannot recover those entries. Adding a memo key function makes future entries type-aware, changes the canonical form, and therefore starts cold; it still cannot recover prior pickle-fallback entries. Keep `prev_type_id` on the moved class for as long as that historical identity is desired. Replacing it later with an ordinary string stable ID intentionally starts a new namespace. Normal decorated-function logic changes can still invalidate memoization. - -`prev_type_id` preserves only this input type's identity within the memo key. Reuse still requires every other call-key component and normal reuse check to remain compatible; moving or renaming the memoized function can therefore start cold independently of the input type. - -Treat that previous canonical `(module, qualname)` as reserved while any active type emits it, and afterward for as long as entries under the historical identity may still be addressed. If a semantically different type later uses the same module and qualified name, matching memo-key payloads can share cached results. - -Assign the value returned by `prev_type_id` directly, or pass it intact as `stable_type_id`. Converting or serializing it to a plain string removes the historical-identity marker. For text-only persistence, store the module and qualified name separately and call `prev_type_id` again after loading. - -If you control the class definition, set `__coco_memo_type_id__` on the class. If you cannot or do not want to edit the class definition, register the same stable type ID with `register_memo_key_function`. - -Use `stable_type_id` without a key function only when CocoIndex already fingerprints the exact runtime type: dataclass instances, Pydantic model instances, instances with `__coco_memo_key__()`, or the class object itself. Subclasses do not inherit this form; give the subclass its own stable type ID unless a registered key function handles it through MRO lookup. - -Rule of thumb: if CocoIndex would otherwise use pickle for your object, `stable_type_id` by itself is not enough; register a key function and stable type ID together. - -Either omit the key function argument or pass `None` when registering only a stable type ID; both forms are equivalent. - -```python -from dataclasses import dataclass -import cocoindex as coco - -@dataclass -class ProductRow: - sku: str - updated_at: int - -PRODUCT_ROW_TYPE_ID = "com.example.ProductRow/v1" - -coco.register_memo_key_function( - ProductRow, - stable_type_id=PRODUCT_ROW_TYPE_ID, -) -``` - -This registration is process-global. Register before the value is fingerprinted: for change-detected context values, before `builder.provide()`; for function arguments, before memoized calls that should use it; and for `deps`, before the `@coco.fn(..., deps=...)` or `@coco.fn.as_async(..., deps=...)` decorator executes. Later calls replace the previous registration for that exact type. - -This registration has two effects: - -- `ProductRow(...)` instances keep their normal dataclass field fingerprint, but the surrounding type namespace uses `PRODUCT_ROW_TYPE_ID`. -- The class object `ProductRow` also uses `PRODUCT_ROW_TYPE_ID` when you pass the class itself (`ProductRow`), not an instance (`ProductRow(...)`). - -If the memoized function takes the class itself and reads class-level details, a stable type ID alone is not enough. Examples include field names, field types/defaults, validators, Pydantic config, class attributes, or version constants. Because `memo_key=` replaces the normal fingerprint for `row_type` in the example below, include both the stable type ID and the class-level details the function reads: - -```python -# Continuing the ProductRow example above. -from dataclasses import fields -import cocoindex as coco - -@coco.fn( - memo=True, - memo_key={ - "row_type": lambda row_type: ( - PRODUCT_ROW_TYPE_ID, - tuple(field.name for field in fields(row_type)), - ) - }, -) -def product_row_fields(row_type: type[ProductRow]) -> tuple[str, ...]: - return tuple(field.name for field in fields(row_type)) - -product_row_fields(ProductRow) -``` - -:::::caution[Class object schema is not fingerprinted automatically] -A stable type ID says the class should keep the same memo namespace after a move or rename. It does not say the class schema is unchanged. If cached results depend on class-level details, include those details with `memo_key=` for that function. If a class change should invalidate every memo that uses this stable type ID, change the stable type ID, for example from `/v1` to `/v2`. -::::: - -Use stable type IDs directly for class objects only when the memoized computation does not depend on class details such as module, qualified name, display name, schema, or state. Those details are intentionally excluded from the class object's fingerprint when a stable type ID is used. - -For a third-party type that would otherwise fall back to pickle, register both a custom key function and a stable type ID on the same owner type: +For types you don't control, register the ID with `register_memo_key_function()`: ```python from some_library import ExternalRecord @@ -234,82 +153,20 @@ coco.register_memo_key_function( ) ``` -Registration is replacement-based, not additive. Each call replaces the full registration for that exact type: - -| Call shape | Registered key/state functions after the call | Registered stable type ID after the call | -| --- | --- | --- | -| `register_memo_key_function(T, key_fn, state_fn=..., stable_type_id="...")` — see [`state_fn` requirements](#register-a-state-function-when-you-dont-control-the-type) | `key_fn` and `state_fn` are installed. | The provided stable type ID is installed. | -| `register_memo_key_function(T, key_fn)` | `key_fn` is installed and any previous `state_fn` is cleared unless a new `state_fn` is supplied. | Any previous registered stable type ID is cleared. | -| `register_memo_key_function(T, stable_type_id="...")` | Any previous key/state functions are cleared. | The provided stable type ID is installed. | - -A `__coco_memo_type_id__` defined on the class is separate from the process-global registration table. If both are set for the same exact type, the class-defined stable type ID takes precedence. The registration's `stable_type_id=` is used only when the selected owner has no non-`None` class-defined ID. - -For a registry-owned stable type ID, if a type needs a key function, state function, and stable type ID, register all of them in the same call. Avoid splitting registration across helper modules where import order can determine which registration wins. - -If a type already has a key or state function and you want to add a stable type ID, update the same call: - -```python -coco.register_memo_key_function( - T, - key_fn, - state_fn=state_fn, - stable_type_id="com.example.T/v1", -) -``` - -Do not add the stable type ID with a separate later call; that later call replaces the registration and clears the previous key/state functions. - -#### Stable type ID without a key function: exact type only - -`__coco_memo_type_id__` and `register_memo_key_function(SomeType, stable_type_id=...)` without a key function apply only to exactly `SomeType`, not subclasses. For example: - -```python -class Parent: - ... - -class Child(Parent): - ... - -coco.register_memo_key_function(Parent, stable_type_id="com.example.Parent/v1") -# Child(...) does not use Parent's stable type ID, because this registration -# has no key function and applies only to exactly Parent. -``` - -#### Stable type ID with a key function: owner selected by MRO - -When `register_memo_key_function(Parent, key_fn, stable_type_id=...)` handles a `Child(Parent)` value through MRO lookup, the selected registered owner (`Parent`) supplies the type namespace. `Parent.__coco_memo_type_id__` supplies that namespace when it is non-`None`; otherwise, the registration's `stable_type_id=...` is used. For example: - -```python -class Parent: - ... - -class Child(Parent): - ... - -coco.register_memo_key_function( - Parent, - lambda value: value.id, - stable_type_id="com.example.Parent/v1", -) -# Child(...) can use Parent's stable type ID when Parent's key function wins -# MRO lookup and handles the Child value. -``` - -Adding a new ordinary string stable type ID does not retroactively migrate memo entries created under the automatic identity or a different string ID. Use `prev_type_id(...)` to retain a previous automatic identity; it keeps those entries directly addressable rather than migrating them. +:::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 +### Class objects as inputs -Class objects are the class values themselves, such as `ProductRow`, not instances such as `ProductRow(...)`. +A class object is the class value itself (such as `ProductRow`), not an instance (such as `ProductRow(...)`). When CocoIndex fingerprints a class object: -- It never calls `__coco_memo_key__` or `__coco_memo_state__` attributes on the class object itself. -- It checks memo key functions explicitly registered on the class object's metaclass or its bases, including `type`. The most specific registered owner wins; a `state_fn` from the same registration participates in state validation, and that owner supplies the stable type ID or module-plus-qualified-name namespace. -- If no metaclass or `type` key function is registered, it uses the class object's own stable type ID when present, otherwise its module plus qualified name. -- A key/state registration on `ProductRow` customizes `ProductRow(...)` instances only. To customize the raw `ProductRow` class globally, register a key function on its metaclass. A stable-ID-only registration on a metaclass or on `type` does not customize all class objects. -- It excludes class schema and class state. If a memoized result depends on fields, field types/defaults, validators, Pydantic config, class attributes, or any other class-level details, pass that data separately with `memo_key=` or change the stable type ID. - -Do not add a `@classmethod` or `@staticmethod` named `__coco_memo_key__` to customize class-object memoization. Class objects ignore it, but instances can still see it and accidentally give every instance the same memo key. +* It uses the class's `__coco_memo_type_id__` when defined, or falls back to its module and qualified name (`module.QualName`). +* It does not call `__coco_memo_key__` or `__coco_memo_state__` methods defined on the class or its metaclass. +* It fingerprints only the class identity, not its schema or attributes. Include class details with `memo_key=` if the function depends on them. +* To customize memoization for class objects globally, register a memo key function on the class's custom metaclass or on `type` (which takes precedence over `__coco_memo_type_id__`). ### Override at the call site with `memo_key=` From 38cd32b2064c375680433627371312cadf729ece Mon Sep 17 00:00:00 2001 From: Max Rong Date: Fri, 28 Aug 2026 06:45:42 +0000 Subject: [PATCH 12/16] refactor(memo): simplify memo registry to type-keyed dictionary - Replace identity-keyed weakref memo registry with direct type-keyed dictionary (dict[type, _MemoTypeRegistry]). - Remove legacy weak-identity test cases (EqMeta, EqNoHashMeta, EqHashMeta). - Remove identity registry lookup monkeypatch tests. --- .../cocoindex/_internal/memo_fingerprint.py | 41 +--- .../tests/internal/test_memo_fingerprint.py | 189 +----------------- 2 files changed, 6 insertions(+), 224 deletions(-) diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index ce5eeb42e..a1667db50 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -16,7 +16,6 @@ import struct import sys import typing -import weakref from . import core from .serde import ( @@ -37,9 +36,7 @@ class _MemoTypeRegistry(typing.NamedTuple): stable_type_id: str | None = None -_memo_type_registry: dict[ - int, tuple[weakref.ReferenceType[type], _MemoTypeRegistry] -] = {} +_memo_type_registry: dict[type, _MemoTypeRegistry] = {} class StateFnEntry(typing.NamedTuple): @@ -165,47 +162,19 @@ def prev_type_id(module: str, qualname: str) -> str: return _PreviousTypeId(module, qualname) -def _remove_memo_type_registry( - type_id: int, dead_ref: weakref.ReferenceType[type] -) -> None: - """Remove ``type_id`` only if ``dead_ref`` is still the stored weakref.""" - entry = _memo_type_registry.get(type_id) - if entry is not None and entry[0] is dead_ref: - _memo_type_registry.pop(type_id, None) - - def _register_memo_type_registry(typ: type, registry: _MemoTypeRegistry) -> None: """Register memo configuration for one exact Python type.""" - type_id = id(typ) - - def _remove_stale_registry(dead_ref: weakref.ReferenceType[type]) -> None: - _remove_memo_type_registry(type_id, dead_ref) - - _memo_type_registry[type_id] = ( - weakref.ref(typ, _remove_stale_registry), - registry, - ) + _memo_type_registry[typ] = registry def _unregister_memo_type_registry(typ: type) -> None: """Best-effort removal of an exact-type memo registration.""" - type_id = id(typ) - entry = _memo_type_registry.get(type_id) - if entry is not None and entry[0]() is typ: - _memo_type_registry.pop(type_id, None) + _memo_type_registry.pop(typ, None) def _registered_memo_type_registry(typ: type) -> _MemoTypeRegistry | None: - """Return the exact type's registration from the identity-keyed table.""" - type_id = id(typ) - entry = _memo_type_registry.get(type_id) - if entry is None: - return None - ref, registry = entry - if ref() is typ: - return registry - _memo_type_registry.pop(type_id, None) - return None + """Return the exact type's registration from the type-keyed table.""" + return _memo_type_registry.get(typ) _MEMO_KEY_ATTR = "__coco_memo_key__" diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index 21f1e54c6..79e3952b0 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -313,6 +313,7 @@ class IdentitySourceEntry: copy.deepcopy(marker), pickle.loads(pickle.dumps(marker)), ): + assert type(variant) is type(marker) class MovedSourceEntry: __coco_memo_type_id__: ClassVar[str] = variant @@ -1115,106 +1116,6 @@ class ChildEntry(BaseEntry): unregister_memo_key_function(BaseEntry) -def test_identity_dispatch_looks_up_each_candidate_registry_once( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class ExactEntry: - pass - - class RegisteredBase: - pass - - class InheritedEntry(RegisteredBase): - pass - - @dataclasses.dataclass - class DataclassEntry: - value: int - - class MemoMeta(type): - pass - - class MetaEntry(metaclass=MemoMeta): - pass - - original_lookup = _memo_fingerprint._registered_memo_type_registry - lookups: list[type] = [] - - def tracked_lookup(typ: type) -> Any: - lookups.append(typ) - return original_lookup(typ) - - def assert_lookups(*expected: type) -> None: - assert lookups == list(expected) - lookups.clear() - - try: - register_memo_key_function(ExactEntry, lambda _entry: "exact") - register_memo_key_function(RegisteredBase, lambda _entry: "inherited") - register_memo_key_function( - DataclassEntry, stable_type_id="test.LookupDataclass/v1" - ) - register_memo_key_function(MemoMeta, lambda _cls: "metaclass") - monkeypatch.setattr( - _memo_fingerprint, "_registered_memo_type_registry", tracked_lookup - ) - - _memo_fingerprint._canonicalize(ExactEntry(), None, []) - assert_lookups(ExactEntry) - - _memo_fingerprint._canonicalize(InheritedEntry(), None, []) - assert_lookups(InheritedEntry, RegisteredBase) - - _memo_fingerprint._canonicalize(DataclassEntry(1), None, []) - assert_lookups(DataclassEntry, object) - - _memo_fingerprint._canonicalize(MetaEntry, None, []) - assert_lookups(MemoMeta) - finally: - unregister_memo_key_function(ExactEntry) - unregister_memo_key_function(RegisteredBase) - unregister_memo_key_function(DataclassEntry) - unregister_memo_key_function(MemoMeta) - - -def test_pydantic_identity_registry_lookup_is_not_repeated( - monkeypatch: pytest.MonkeyPatch, -) -> None: - try: - from pydantic import BaseModel - except ImportError: - pytest.skip("pydantic not installed") - return - - class Model(BaseModel): - value: int - - original_lookup = _memo_fingerprint._registered_memo_type_registry - lookups: list[type] = [] - - def tracked_lookup(typ: type) -> Any: - lookups.append(typ) - return original_lookup(typ) - - try: - register_memo_key_function(Model, stable_type_id="test.LookupPydantic/v1") - monkeypatch.setattr( - _memo_fingerprint, "_registered_memo_type_registry", tracked_lookup - ) - canonical = _memo_fingerprint._canonicalize(Model(value=1), None, []) - assert isinstance(canonical, tuple) - - assert canonical[:3] == ( - "pydantic", - ("__coco_memo_type_id__", "test.LookupPydantic/v1"), - None, - ) - assert sum(owner is Model for owner in lookups) == 1 - assert len({id(owner) for owner in lookups}) == len(lookups) - finally: - unregister_memo_key_function(Model) - - def test_combined_registration_uses_stable_type_id_and_collects_state_fn() -> None: class OldEntry: def __init__(self, value: object) -> None: @@ -1471,94 +1372,6 @@ def __coco_memo_key__(self) -> object: unregister_memo_key_function(NewEntry) -def test_registered_stable_type_id_is_identity_exact_for_equal_metaclasses() -> None: - class EqMeta(type): - def __eq__(cls, other: object) -> bool: - return isinstance(other, EqMeta) - - def __hash__(cls) -> int: - return 1 - - class A(metaclass=EqMeta): - def __coco_memo_key__(self) -> object: - return ("same",) - - class B(metaclass=EqMeta): - def __coco_memo_key__(self) -> object: - return ("same",) - - assert fingerprint_call(_dummy_fn, (A(),), {}, []) != fingerprint_call( - _dummy_fn, (B(),), {}, [] - ) - - try: - register_memo_key_function(A, stable_type_id="test.EqualityMetaA/v1") - assert fingerprint_call(_dummy_fn, (A(),), {}, []) != fingerprint_call( - _dummy_fn, (B(),), {}, [] - ) - finally: - unregister_memo_key_function(A) - - -def test_unregister_memo_key_function_handles_unhashable_stable_type_id_only() -> None: - class EqNoHashMeta(type): - def __eq__(cls, other: object) -> bool: - return cls is other - - class OldEntry(metaclass=EqNoHashMeta): - def __init__(self, value: object) -> None: - self.value = value - - def __coco_memo_key__(self) -> object: - return ("entry", self.value) - - class NewEntry(metaclass=EqNoHashMeta): - def __init__(self, value: object) -> None: - self.value = value - - def __coco_memo_key__(self) -> object: - return ("entry", self.value) - - try: - register_memo_key_function(OldEntry, stable_type_id="test.UnhashableMeta/v1") - register_memo_key_function(NewEntry, stable_type_id="test.UnhashableMeta/v1") - assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( - _dummy_fn, (NewEntry(1),), {}, [] - ) - unregister_memo_key_function(OldEntry) - assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) != fingerprint_call( - _dummy_fn, (NewEntry(1),), {}, [] - ) - finally: - unregister_memo_key_function(OldEntry) - unregister_memo_key_function(NewEntry) - - -def test_unregister_memo_key_function_is_identity_exact_for_equal_metaclasses() -> None: - class EqHashMeta(type): - def __eq__(cls, other: object) -> bool: - return isinstance(other, EqHashMeta) - - def __hash__(cls) -> int: - return 1 - - class StableOnly(metaclass=EqHashMeta): - pass - - class Registered(metaclass=EqHashMeta): - def __getstate__(self) -> object: - raise TypeError("registered test object is not picklable") - - try: - register_memo_key_function(Registered, lambda entry: ("registered",)) - register_memo_key_function(StableOnly, stable_type_id="test.EqualUnregister/v1") - unregister_memo_key_function(StableOnly) - fingerprint_call(_dummy_fn, (Registered(),), {}, []) - finally: - unregister_memo_key_function(StableOnly) - unregister_memo_key_function(Registered) - - def test_stable_type_id_exact_type() -> None: class OldEntry: __coco_memo_type_id__ = "test.DirectEntry/v1" From 9ea69491bbb3a33c288f11fa585d12f4819052f6 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 29 Aug 2026 15:58:47 +0000 Subject: [PATCH 13/16] fix(memo): enforce string type checks for stable type IDs - Narrow stable_type_id check to isinstance(..., str) in _type_identity_parts. - Validate that stable_type_id is a str when provided in register_memo_key_function. - Add regression tests covering invalid stable_type_id in registration and class attributes. --- python/cocoindex/_internal/memo_fingerprint.py | 7 ++++++- python/tests/internal/test_memo_fingerprint.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index a1667db50..d4aa8c1a2 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -202,7 +202,7 @@ def _type_identity_parts( if isinstance(stable_type_id, _PreviousTypeId): return stable_type_id._identity_parts - if stable_type_id is not None: + if isinstance(stable_type_id, str): return (("__coco_memo_type_id__", stable_type_id), None) return (canonical_module_name(typ), getattr(typ, "__qualname__", None)) @@ -404,6 +404,11 @@ def register_memo_key_function( "register_memo_key_function() 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_key_function() stable_type_id must be a str, " + f"got {type(stable_type_id).__name__}" + ) _register_memo_type_registry( typ, diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index 79e3952b0..fad3e7345 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -348,6 +348,14 @@ class OrdinaryStringSourceEntry: None, ) + class NonStringSourceEntry: + __coco_memo_type_id__ = 12345 + + assert _memo_fingerprint._type_identity_parts(NonStringSourceEntry, None) == ( + _memo_fingerprint.canonical_module_name(NonStringSourceEntry), + NonStringSourceEntry.__qualname__, + ) + def test_pydantic_stable_type_id_allows_renamed_model_reuse() -> None: try: @@ -1279,6 +1287,11 @@ def __coco_memo_key__(self) -> object: {"state_fn": object()}, "state_fn must be callable", ), + ( + (), + {"stable_type_id": 123}, + "stable_type_id must be a str", + ), ], ) def test_register_memo_key_function_rejects_invalid_forms( From c64b3a7aba3f8ceb6246a4bb428b64b6e4b19eb8 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sat, 29 Aug 2026 22:29:00 +0000 Subject: [PATCH 14/16] feat(memo): inherit declared stable type IDs via getattr with MRO registration fallback - Switch stable type ID lookup to getattr(typ, "__coco_memo_type_id__", None) so declared stable type IDs inherit down subclasses. - Allow registered stable type IDs to apply to subclasses across MRO when no declared ID is present. - Enforce declared-beats-registered precedence (declared ID on base beats registered ID on subclass). - Let a subclass override an inherited registered stable ID by declaring __coco_memo_type_id__ or registering its own; instance identity resolves on the runtime type with the automatic module/qualname fallback anchored to the key function's MRO owner via fallback_owner, preserving prev_type_id migration semantics. - Document stable type ID inheritance, precedence, MRO resolution, full-replace semantics, and prev_type_id base propagation. - Add regression tests covering inheritance, MRO resolution, declared precedence, subclass overrides, and class objects. --- .../docs/advanced_topics/memoization_keys.mdx | 12 +- .../cocoindex/_internal/memo_fingerprint.py | 53 +++- .../tests/internal/test_memo_fingerprint.py | 251 +++++++++++++----- 3 files changed, 234 insertions(+), 82 deletions(-) diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index df0aba22b..6cda13056 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -140,6 +140,8 @@ class SourceEntry: # Moved to new_package.models 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_key_function()`: ```python @@ -153,6 +155,12 @@ coco.register_memo_key_function( ) ``` +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 propagates to subclasses. A subclass can override the ID by declaring its own. +- **Precedence**: A declared `__coco_memo_type_id__` (own or inherited) takes precedence over any registered stable type ID. +- **MRO registration**: Registered stable type IDs also apply to subclasses across the MRO: the most specific registered type wins. A subclass can override an inherited registered ID by declaring or registering its own. + :::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. ::: @@ -163,10 +171,10 @@ A class object is the class value itself (such as `ProductRow`), not an instance When CocoIndex fingerprints a class object: -* It uses the class's `__coco_memo_type_id__` when defined, or falls back to its module and qualified name (`module.QualName`). +* It uses the class's stable type ID (declared with `__coco_memo_type_id__` or registered with `register_memo_key_function()`) when defined, or falls back to its module and qualified name (`module.QualName`). * It does not call `__coco_memo_key__` or `__coco_memo_state__` methods defined on the class or its metaclass. * It fingerprints only the class identity, not its schema or attributes. Include class details with `memo_key=` if the function depends on them. -* To customize memoization for class objects globally, register a memo key function on the class's custom metaclass or on `type` (which takes precedence over `__coco_memo_type_id__`). +* To customize memoization for class objects globally, register a memo key function on the class's custom metaclass or on `type` (which takes precedence over stable type IDs). ### Override at the call site with `memo_key=` diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index d4aa8c1a2..c0d43da84 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -188,6 +188,7 @@ def _registered_memo_type_registry(typ: type) -> _MemoTypeRegistry | None: def _type_identity_parts( typ: type, registry: _MemoTypeRegistry | None, + fallback_owner: type | None = None, ) -> tuple[Fingerprintable, Fingerprintable]: """Return stable type ID or module+qualname type identity parts. @@ -195,16 +196,34 @@ def _type_identity_parts( 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 = typ.__dict__.get("__coco_memo_type_id__") - if stable_type_id is None and registry is not None: - stable_type_id = registry.stable_type_id - + 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__: + reg = ( + registry + if 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) - return (canonical_module_name(typ), getattr(typ, "__qualname__", 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( @@ -222,9 +241,23 @@ def _canonicalize_registered_memo_key( 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) + 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, - *_type_identity_parts(owner, registry), + *identity, _canonicalize(key, state, state_methods), ) @@ -366,10 +399,8 @@ def register_memo_key_function( ) -> None: """Register a memo key function and/or stable type ID for a type. - Key-function resolution is MRO-aware: the most specific registered base - type wins. Stable type IDs registered without a key function apply to the - exact type only; stable type IDs registered with a key function identify - that selected owner type. Each call replaces the full registration for + 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. @@ -446,7 +477,7 @@ def _raise_not_memo_keyable(obj: object) -> typing.NoReturn: def unregister_memo_key_function(typ: type) -> None: - """Remove registered memo key function and stable type ID (best-effort).""" + """Remove registered memo key function and stable type ID.""" if isinstance(typ, type): _unregister_memo_type_registry(typ) diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index fad3e7345..56448d3b5 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -815,16 +815,22 @@ class Entry: unregister_memo_key_function(type) -def test_raw_class_object_stable_type_id_is_exact_type() -> None: +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 - assert fingerprint_call(_dummy_fn, (Parent,), {}, []) != fingerprint_call( + 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_key_function_registers_key_function_and_stable_type_id_for_owner_base() -> ( @@ -919,7 +925,7 @@ def __coco_memo_key__(self) -> object: try: register_memo_key_function(OldEntry, stable_type_id="test.RegisteredEntry/v1") - register_memo_key_function(NewEntry, stable_type_id="test.RegisteredEntry/v1") + register_memo_key_function(NewEntry, None, stable_type_id="test.RegisteredEntry/v1") assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) == fingerprint_call( _dummy_fn, (NewEntry(1),), {}, [] ) @@ -931,7 +937,7 @@ def __coco_memo_key__(self) -> object: unregister_memo_key_function(NewEntry) -def test_stable_type_id_only_registration_is_exact_for_subclasses() -> None: +def test_stable_type_id_only_registration_propagates_to_subclasses_across_mro() -> None: class Parent: def __init__(self, value: object) -> None: self.value = value @@ -942,6 +948,9 @@ def __coco_memo_key__(self) -> object: class Child(Parent): pass + class OverridingChild(Parent): + pass + class SameStableTypeId: def __init__(self, value: object) -> None: self.value = value @@ -951,26 +960,116 @@ def __coco_memo_key__(self) -> object: try: register_memo_key_function( - Parent, stable_type_id="test.RegisteredExactParent/v1" + Parent, stable_type_id="test.RegisteredParent/v1" + ) + register_memo_key_function( + OverridingChild, stable_type_id="test.RegisteredOverridingChild/v1" ) register_memo_key_function( - SameStableTypeId, stable_type_id="test.RegisteredExactParent/v1" + 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, (SameStableTypeId(1),), {}, [] - ) - assert fingerprint_call(_dummy_fn, (Parent(1),), {}, []) != fingerprint_call( _dummy_fn, (Child(1),), {}, [] ) - assert fingerprint_call(_dummy_fn, (Parent,), {}, []) != fingerprint_call( + 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_key_function(Parent) + unregister_memo_key_function(OverridingChild) unregister_memo_key_function(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_key_function(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_key_function(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: @@ -995,6 +1094,78 @@ class Child(Parent): unregister_memo_key_function(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_key_function( + 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_key_function(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_key_function( + Base, + lambda entry: ("base", entry.value), + stable_type_id="test.BaseRegisteredOverride/v2", + ) + register_memo_key_function(Sub, stable_type_id="test.SubRegisteredOverride/v2") + register_memo_key_function(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_key_function(TwinSub) + unregister_memo_key_function(Sub) + unregister_memo_key_function(Base) + + def test_stable_type_id_only_registration_replaces_key_and_state_functions() -> None: @dataclasses.dataclass class Entry: @@ -1255,27 +1426,6 @@ class SameStableTypeId: unregister_memo_key_function(SameStableTypeId) -def test_register_memo_key_function_accepts_explicit_none_key_function() -> None: - class Entry: - def __coco_memo_key__(self) -> object: - return ("entry",) - - class SameStableTypeId: - def __coco_memo_key__(self) -> object: - return ("entry",) - - try: - register_memo_key_function(Entry, None, stable_type_id="test.ExplicitNone/v1") - register_memo_key_function( - SameStableTypeId, stable_type_id="test.ExplicitNone/v1" - ) - assert fingerprint_call(_dummy_fn, (Entry(),), {}, []) == fingerprint_call( - _dummy_fn, (SameStableTypeId(),), {}, [] - ) - finally: - unregister_memo_key_function(Entry) - unregister_memo_key_function(SameStableTypeId) - @pytest.mark.parametrize( ("args", "kwargs", "match"), @@ -1385,49 +1535,12 @@ def __coco_memo_key__(self) -> object: unregister_memo_key_function(NewEntry) -def test_stable_type_id_exact_type() -> None: - class OldEntry: - __coco_memo_type_id__ = "test.DirectEntry/v1" - - def __init__(self, value: object) -> None: - self.value = value - - def __coco_memo_key__(self) -> object: - return ("entry", self.value) - - class NewEntry: - __coco_memo_type_id__ = "test.DirectEntry/v1" - - def __init__(self, value: object) -> None: - self.value = value - - def __coco_memo_key__(self) -> object: - return ("entry", self.value) - - class Parent: - __coco_memo_type_id__ = "test.Parent/v1" - - def __coco_memo_key__(self) -> object: - return ("same", 1) - - class Child(Parent): - pass - +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) - - 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),), {}, [] - ) - assert fingerprint_call(_dummy_fn, (Parent(),), {}, []) != fingerprint_call( - _dummy_fn, (Child(),), {}, [] - ) none_id_canonical = _memo_fingerprint._canonicalize(NoneId(), None, []) assert isinstance(none_id_canonical, tuple) assert none_id_canonical[:3] == ( From a4e2cb47ea9cb87c4c202e620072a09872d73930 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sun, 30 Aug 2026 00:40:28 +0000 Subject: [PATCH 15/16] refactor(memo): introduce register_memo_type and preserve register_memo_key_function compatibility - Introduce register_memo_type with full-replace semantics as the primary registration API. - Restore register_memo_key_function as a backward-compatible shortcut preserving existing stable type IDs. - Implement unregister_memo_type and unregister_memo_key_function as internal test helpers. - Migrate unit tests to register_memo_type and clean up duplicate test boilerplate. --- .../cocoindex/_internal/memo_fingerprint.py | 54 ++- .../tests/internal/test_memo_fingerprint.py | 395 +++++++++--------- 2 files changed, 249 insertions(+), 200 deletions(-) diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index c0d43da84..e011eaa8d 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -209,7 +209,7 @@ def _type_identity_parts( for owner in typ.__mro__: reg = ( registry - if owner is typ + 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: @@ -372,7 +372,7 @@ def __coco_memo_key__(self) -> typing.NoReturn: @typing.overload -def register_memo_key_function( +def register_memo_type( typ: type, key_fn: _KeyFn, *, @@ -382,7 +382,7 @@ def register_memo_key_function( @typing.overload -def register_memo_key_function( +def register_memo_type( typ: type, key_fn: None = None, *, @@ -390,14 +390,14 @@ def register_memo_key_function( ) -> None: ... -def register_memo_key_function( +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 and/or stable type ID for a type. + """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 @@ -413,31 +413,31 @@ def register_memo_key_function( if not isinstance(typ, type): raise TypeError( - "register_memo_key_function() expects typ to be a type, " + "register_memo_type() expects typ to be a type, " f"got {type(typ).__name__}" ) if key_fn is None: if state_fn is not None: raise TypeError( - "register_memo_key_function() state_fn requires a memo key function" + "register_memo_type() state_fn requires a memo key function" ) if stable_type_id is None: raise TypeError( - "register_memo_key_function() requires a key_fn or stable_type_id" + "register_memo_type() requires a key_fn or stable_type_id" ) elif not callable(key_fn): raise TypeError( - "register_memo_key_function() key_fn must be callable, " + "register_memo_type() key_fn must be callable, " f"got {type(key_fn).__name__}" ) if state_fn is not None and not callable(state_fn): raise TypeError( - "register_memo_key_function() state_fn must be callable, " + "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_key_function() stable_type_id must be a str, " + "register_memo_type() stable_type_id must be a str, " f"got {type(stable_type_id).__name__}" ) @@ -451,6 +451,18 @@ def register_memo_key_function( ) +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) + + def register_not_memo_keyable(typ: type) -> None: """Register a type as not memo-keyable. @@ -476,13 +488,27 @@ def _raise_not_memo_keyable(obj: object) -> typing.NoReturn: _register_memo_type_registry(typ, _MemoTypeRegistry(key_fn=_raise_not_memo_keyable)) -def unregister_memo_key_function(typ: type) -> None: - """Remove registered memo key function and stable type ID.""" +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 registered memo key function, preserving any registered stable type ID.""" + + 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, ...]: """Return a totally-ordered key for canonical values. @@ -688,6 +714,8 @@ def fingerprint_call( "memo_fingerprint", "prev_type_id", "register_memo_key_function", + "register_memo_type", "register_not_memo_keyable", "unregister_memo_key_function", + "unregister_memo_type", ] diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index 56448d3b5..5473a4124 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -12,8 +12,10 @@ 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 @@ -327,7 +329,7 @@ class RegisteredSourceEntry: pass try: - register_memo_key_function(RegisteredSourceEntry, stable_type_id=marker) + register_memo_type(RegisteredSourceEntry, stable_type_id=marker) registry = _memo_fingerprint._registered_memo_type_registry( RegisteredSourceEntry ) @@ -336,7 +338,7 @@ class RegisteredSourceEntry: RegisteredSourceEntry, registry ) == (module, qualname) finally: - unregister_memo_key_function(RegisteredSourceEntry) + unregister_memo_type(RegisteredSourceEntry) ordinary_marker = str(marker) @@ -427,13 +429,13 @@ class ChangedEntry: pass try: - register_memo_key_function( + register_memo_type( OldEntry, stable_type_id="test.RegisteredRawClass/v1" ) - register_memo_key_function( + register_memo_type( NewEntry, stable_type_id="test.RegisteredRawClass/v1" ) - register_memo_key_function( + register_memo_type( ChangedEntry, stable_type_id="test.RegisteredRawClass/v2" ) @@ -444,9 +446,9 @@ class ChangedEntry: _dummy_fn, (ChangedEntry,), {}, [] ) finally: - unregister_memo_key_function(OldEntry) - unregister_memo_key_function(NewEntry) - unregister_memo_key_function(ChangedEntry) + unregister_memo_type(OldEntry) + unregister_memo_type(NewEntry) + unregister_memo_type(ChangedEntry) def test_hook_memo_key_fragment_preserves_parent_cycle() -> None: @@ -623,7 +625,7 @@ def metaclass_state(cls: type, prev_state: object) -> MemoStateOutcome: stable_type_id = "test.RawClassRegisteredMetaOwner/v1" try: - register_memo_key_function( + register_memo_type( MemoMeta, metaclass_key, state_fn=metaclass_state, @@ -664,7 +666,7 @@ def metaclass_state(cls: type, prev_state: object) -> MemoStateOutcome: (NewEntry, "reusable"), ] finally: - unregister_memo_key_function(MemoMeta) + unregister_memo_type(MemoMeta) def test_prev_type_id_reuses_registered_metaclass_owner_identity() -> None: @@ -689,8 +691,8 @@ def metaclass_key(cls: type) -> object: return "source-class" try: - register_memo_key_function(OldMemoMeta, metaclass_key) - register_memo_key_function( + register_memo_type(OldMemoMeta, metaclass_key) + register_memo_type( MovedMemoMeta, metaclass_key, stable_type_id=coco.prev_type_id( @@ -702,8 +704,8 @@ def metaclass_key(cls: type) -> object: _dummy_fn, (MovedEntry,), {}, [] ) finally: - unregister_memo_key_function(OldMemoMeta) - unregister_memo_key_function(MovedMemoMeta) + unregister_memo_type(OldMemoMeta) + unregister_memo_type(MovedMemoMeta) def test_raw_class_object_honors_registered_type_memo_key_and_state() -> None: @@ -737,7 +739,7 @@ def type_state(cls: type, prev_state: object) -> MemoStateOutcome: stable_type_id = "test.RawClassRegisteredTypeOwner/v1" try: - register_memo_key_function( + register_memo_type( type, type_key, state_fn=type_state, @@ -778,7 +780,7 @@ def type_state(cls: type, prev_state: object) -> MemoStateOutcome: (NewEntry, "reusable"), ] finally: - unregister_memo_key_function(type) + unregister_memo_type(type) def test_raw_class_object_ignores_registered_object_memo_key() -> None: @@ -794,11 +796,11 @@ def object_key(obj: object) -> object: return "object memo key ran" try: - register_memo_key_function(object, object_key) + register_memo_type(object, object_key) assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == expected assert object_key_calls == [] finally: - unregister_memo_key_function(object) + unregister_memo_type(object) def test_raw_class_object_ignores_registered_type_stable_type_id() -> None: @@ -809,10 +811,10 @@ class Entry: original = fingerprint_call(_dummy_fn, (Entry,), {}, []) try: - register_memo_key_function(type, stable_type_id="test.RegisteredType/v1") + register_memo_type(type, stable_type_id="test.RegisteredType/v1") assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == original finally: - unregister_memo_key_function(type) + unregister_memo_type(type) def test_raw_class_object_stable_type_id_is_inherited_by_subclasses() -> None: @@ -833,7 +835,7 @@ class OverridingChild(Parent): ) -def test_register_memo_key_function_registers_key_function_and_stable_type_id_for_owner_base() -> ( +def test_register_memo_type_registers_key_function_and_stable_type_id_for_owner_base() -> ( None ): class Base: @@ -847,7 +849,7 @@ class ChildB(Base): pass try: - register_memo_key_function( + register_memo_type( Base, lambda entry: ("base", entry.value), stable_type_id="test.RegisteredBase/v1", @@ -859,7 +861,7 @@ class ChildB(Base): _dummy_fn, (ChildB(2),), {}, [] ) finally: - unregister_memo_key_function(Base) + unregister_memo_type(Base) def test_prev_type_id_reuses_registered_mro_owner_identity() -> None: @@ -886,8 +888,8 @@ def base_key(entry: OldBase | MovedBase) -> object: return ("base", entry.value) try: - register_memo_key_function(OldBase, base_key) - register_memo_key_function( + register_memo_type(OldBase, base_key) + register_memo_type( MovedBase, base_key, stable_type_id=coco.prev_type_id( @@ -902,11 +904,11 @@ def base_key(entry: OldBase | MovedBase) -> object: _dummy_fn, (MovedChild(2),), {}, [] ) finally: - unregister_memo_key_function(OldBase) - unregister_memo_key_function(MovedBase) + unregister_memo_type(OldBase) + unregister_memo_type(MovedBase) -def test_register_memo_key_function_registers_stable_type_id_without_key_function() -> ( +def test_register_memo_type_registers_stable_type_id_without_key_function() -> ( None ): class OldEntry: @@ -924,8 +926,8 @@ def __coco_memo_key__(self) -> object: return ("entry", self.value) try: - register_memo_key_function(OldEntry, stable_type_id="test.RegisteredEntry/v1") - register_memo_key_function(NewEntry, None, stable_type_id="test.RegisteredEntry/v1") + 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),), {}, [] ) @@ -933,8 +935,8 @@ def __coco_memo_key__(self) -> object: _dummy_fn, (NewEntry(2),), {}, [] ) finally: - unregister_memo_key_function(OldEntry) - unregister_memo_key_function(NewEntry) + unregister_memo_type(OldEntry) + unregister_memo_type(NewEntry) def test_stable_type_id_only_registration_propagates_to_subclasses_across_mro() -> None: @@ -959,13 +961,13 @@ def __coco_memo_key__(self) -> object: return ("entry", self.value) try: - register_memo_key_function( + register_memo_type( Parent, stable_type_id="test.RegisteredParent/v1" ) - register_memo_key_function( + register_memo_type( OverridingChild, stable_type_id="test.RegisteredOverridingChild/v1" ) - register_memo_key_function( + register_memo_type( SameStableTypeId, stable_type_id="test.RegisteredParent/v1" ) @@ -984,9 +986,9 @@ def __coco_memo_key__(self) -> object: _dummy_fn, (OverridingChild(1),), {}, [] ) finally: - unregister_memo_key_function(Parent) - unregister_memo_key_function(OverridingChild) - unregister_memo_key_function(SameStableTypeId) + 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: @@ -1007,7 +1009,7 @@ class OverridingSub(Base): try: # Register an ID on Sub — declared (inherited from Base) must take precedence: - register_memo_key_function(Sub, stable_type_id="test.RegisteredSub/v1") + 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"), @@ -1030,7 +1032,7 @@ class OverridingSub(Base): _dummy_fn, (OverridingSub(1),), {}, [] ) finally: - unregister_memo_key_function(Sub) + unregister_memo_type(Sub) def test_prev_type_id_on_base_class_propagates_to_subclasses() -> None: @@ -1070,6 +1072,7 @@ class MovedSub(MovedBase): 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: @@ -1080,8 +1083,8 @@ class Child(Parent): pass try: - register_memo_key_function(Parent, lambda entry: ("parent", entry.value)) - register_memo_key_function(Child, stable_type_id="test.RegisteredExactChild/v1") + 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"),), {}, []) @@ -1090,8 +1093,8 @@ class Child(Parent): fingerprint_call(_dummy_fn, (Child(2, "a"),), {}, []) ) finally: - unregister_memo_key_function(Child) - unregister_memo_key_function(Parent) + unregister_memo_type(Child) + unregister_memo_type(Parent) def test_subclass_declared_stable_type_id_overrides_registered_base_identity() -> None: @@ -1107,7 +1110,7 @@ class TwinSub(Base): __coco_memo_type_id__: ClassVar[str] = "test.SubDeclaredOverride/v1" try: - register_memo_key_function( + register_memo_type( Base, lambda entry: ("base", entry.value), stable_type_id="test.BaseRegisteredOverride/v1", @@ -1125,7 +1128,7 @@ class TwinSub(Base): _dummy_fn, (TwinSub(1, "a"),), {}, [] ) finally: - unregister_memo_key_function(Base) + unregister_memo_type(Base) def test_subclass_registered_stable_type_id_overrides_registered_base_identity() -> None: @@ -1141,13 +1144,13 @@ class TwinSub(Base): pass try: - register_memo_key_function( + register_memo_type( Base, lambda entry: ("base", entry.value), stable_type_id="test.BaseRegisteredOverride/v2", ) - register_memo_key_function(Sub, stable_type_id="test.SubRegisteredOverride/v2") - register_memo_key_function(TwinSub, stable_type_id="test.SubRegisteredOverride/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( @@ -1161,70 +1164,94 @@ class TwinSub(Base): _dummy_fn, (TwinSub(1, "a"),), {}, [] ) finally: - unregister_memo_key_function(TwinSub) - unregister_memo_key_function(Sub) - unregister_memo_key_function(Base) + unregister_memo_type(TwinSub) + unregister_memo_type(Sub) + unregister_memo_type(Base) -def test_stable_type_id_only_registration_replaces_key_and_state_functions() -> None: +def test_register_memo_type_full_replacement_lifecycle() -> None: @dataclasses.dataclass class Entry: value: int - marker: str + ignored: str + + @dataclasses.dataclass + class TwinV1Shook: + value: int + ignored: str - def state_fn(obj: Entry, prev_state: object) -> MemoStateOutcome: - return MemoStateOutcome(state=prev_state, memo_valid=True) + @dataclasses.dataclass + class TwinV2Hook: + value: int + ignored: str + + @dataclasses.dataclass + class TwinV2Dataclass: + value: int + ignored: str + + id_v1 = "test.Lifecycle/v1" + id_v2 = "test.Lifecycle/v2" + + def state_fn(obj: Any, prev_state: object) -> MemoStateOutcome: + return MemoStateOutcome(state=("state", obj.value, prev_state), memo_valid=True) try: - register_memo_key_function( - Entry, - lambda entry: ("constant",), + register_memo_type( + TwinV1Shook, + lambda e: ("key", e.value), state_fn=state_fn, + stable_type_id=id_v1, ) - methods: list[Any] = [] - assert fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) == ( - fingerprint_call(_dummy_fn, (Entry(2, "b"),), {}, []) + register_memo_type( + TwinV2Hook, lambda e: ("key", e.value), stable_type_id=id_v2 ) - assert len(methods) == 1 + register_memo_type(TwinV2Dataclass, stable_type_id=id_v2) - register_memo_key_function( + # 1. Full registration (key_fn, state_fn, stable_id_v1): + register_memo_type( Entry, - stable_type_id="test.ReplaceKeyWithStable/v1", - ) - methods = [] - assert fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) != ( - fingerprint_call(_dummy_fn, (Entry(2, "b"),), {}, []) - ) - assert methods == [] - finally: - unregister_memo_key_function(Entry) - - -def test_key_only_registration_replaces_stable_type_id() -> None: - class Entry: - def __init__(self, value: object) -> None: - self.value = value - - class SameStableTypeId: - def __init__(self, value: object) -> None: - self.value = value - - try: - register_memo_key_function(Entry, stable_type_id="test.ReplaceStableWithKey/v1") - register_memo_key_function( - SameStableTypeId, - lambda entry: ("entry", entry.value), - stable_type_id="test.ReplaceStableWithKey/v1", + lambda e: ("key", e.value), + state_fn=state_fn, + stable_type_id=id_v1, ) + methods: list[Any] = [] + fp1 = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) + assert fp1 == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) # key_fn active (ignores 'ignored') + assert fp1 == fingerprint_call(_dummy_fn, (TwinV1Shook(1, "a"),), {}, []) # identity is id_v1 and tag is shook + assert len(methods) == 1 # state_fn active - register_memo_key_function(Entry, lambda entry: ("entry", entry.value)) - assert fingerprint_call(_dummy_fn, (Entry(1),), {}, []) != fingerprint_call( - _dummy_fn, (SameStableTypeId(1),), {}, [] + # 2. Replace with (key_fn, stable_id_v2): clears state_fn, updates stable_id to id_v2, keeps key_fn + register_memo_type( + Entry, + lambda e: ("key", e.value), + stable_type_id=id_v2, ) + methods = [] + fp2 = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) + assert len(methods) == 0 # state_fn cleared (tag is now hook) + assert fp2 != fp1 # stable_id changed and state_fn cleared + assert fp2 == fingerprint_call(_dummy_fn, (TwinV2Hook(1, "a"),), {}, []) # matches TwinV2Hook + assert fp2 == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) # key_fn still active + + # 3. Replace with stable_id_v2 only: clears key_fn while keeping stable_id_v2 + register_memo_type(Entry, stable_type_id=id_v2) + fp3 = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) + assert fp3 != fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) # key_fn cleared (dataclass default participates) + assert fp3 == fingerprint_call(_dummy_fn, (TwinV2Dataclass(1, "a"),), {}, []) # stable_id_v2 preserved + + # 4. Replace with key_fn only: clears stable_id (identity falls back to module.qualname) + register_memo_type(Entry, lambda e: ("key", e.value)) + fp4 = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) + assert fp4 == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) # key_fn active again + assert fp4 != fingerprint_call(_dummy_fn, (TwinV1Shook(1, "a"),), {}, []) # stable_id cleared + assert fp4 != fingerprint_call(_dummy_fn, (TwinV2Hook(1, "a"),), {}, []) + assert fp4 != fingerprint_call(_dummy_fn, (TwinV2Dataclass(1, "a"),), {}, []) finally: - unregister_memo_key_function(Entry) - unregister_memo_key_function(SameStableTypeId) - + unregister_memo_type(Entry) + unregister_memo_type(TwinV1Shook) + unregister_memo_type(TwinV2Hook) + unregister_memo_type(TwinV2Dataclass) def test_intrinsic_hooks_and_declared_stable_id_beat_registration() -> None: declared_type_id = "test.DeclaredIntrinsic/v1" @@ -1245,7 +1272,7 @@ def registered_state(_entry: Entry, _prev_state: object) -> MemoStateOutcome: raise AssertionError("intrinsic state must beat the registered state") try: - register_memo_key_function( + register_memo_type( Entry, registered_key, state_fn=registered_state, @@ -1264,7 +1291,7 @@ def registered_state(_entry: Entry, _prev_state: object) -> MemoStateOutcome: state="previous", memo_valid=True ) finally: - unregister_memo_key_function(Entry) + unregister_memo_type(Entry) def test_declared_stable_id_beats_registration_for_selected_key_owner() -> None: @@ -1280,7 +1307,7 @@ class ChildEntry(BaseEntry): pass try: - register_memo_key_function( + register_memo_type( BaseEntry, lambda entry: ("registered", entry.value), stable_type_id="test.RegisteredBaseOwner/v1", @@ -1292,7 +1319,7 @@ class ChildEntry(BaseEntry): ("seq", ("registered", 1)), ) finally: - unregister_memo_key_function(BaseEntry) + unregister_memo_type(BaseEntry) def test_combined_registration_uses_stable_type_id_and_collects_state_fn() -> None: @@ -1313,13 +1340,13 @@ def state_fn(entry: Any, prev_state: object) -> MemoStateOutcome: ) try: - register_memo_key_function( + register_memo_type( OldEntry, key_fn, state_fn=state_fn, stable_type_id="test.CombinedStateStable/v1", ) - register_memo_key_function( + register_memo_type( NewEntry, key_fn, state_fn=state_fn, @@ -1337,65 +1364,9 @@ def state_fn(entry: Any, prev_state: object) -> MemoStateOutcome: assert len(methods) == 1 assert methods[0].call("prev").state == ("state", 1, "prev") finally: - unregister_memo_key_function(OldEntry) - unregister_memo_key_function(NewEntry) - - -def test_register_memo_key_function_full_registration_replaces_previous_full_registration() -> ( - None -): - class Entry: - def __init__(self, value: object, marker: object) -> None: - self.value = value - self.marker = marker + unregister_memo_type(OldEntry) + unregister_memo_type(NewEntry) - def old_state_fn(entry: Any, prev_state: object) -> MemoStateOutcome: - return MemoStateOutcome(state=("old", entry.value, prev_state), memo_valid=True) - - def new_state_fn(entry: Any, prev_state: object) -> MemoStateOutcome: - return MemoStateOutcome( - state=("new", entry.marker, prev_state), memo_valid=True - ) - - try: - register_memo_key_function( - Entry, - lambda entry: ("old", entry.value), - state_fn=old_state_fn, - stable_type_id="test.ReplaceFull/old", - ) - first_fingerprint = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) - assert first_fingerprint == fingerprint_call( - _dummy_fn, (Entry(1, "b"),), {}, [] - ) - assert first_fingerprint != fingerprint_call( - _dummy_fn, (Entry(2, "a"),), {}, [] - ) - methods: list[Any] = [] - fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) - assert len(methods) == 1 - assert methods[0].call("prev").state == ("old", 1, "prev") - - register_memo_key_function( - Entry, - lambda entry: ("new", entry.marker), - state_fn=new_state_fn, - stable_type_id="test.ReplaceFull/new", - ) - second_fingerprint = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) - assert first_fingerprint != second_fingerprint - assert second_fingerprint == fingerprint_call( - _dummy_fn, (Entry(2, "a"),), {}, [] - ) - assert second_fingerprint != fingerprint_call( - _dummy_fn, (Entry(1, "b"),), {}, [] - ) - methods = [] - fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) - assert len(methods) == 1 - assert methods[0].call("prev").state == ("new", "a", "prev") - finally: - unregister_memo_key_function(Entry) def test_register_not_memo_keyable_replaces_stable_type_id_for_class_objects() -> None: @@ -1406,10 +1377,10 @@ class SameStableTypeId: pass try: - register_memo_key_function( + register_memo_type( Entry, stable_type_id="test.NotMemoKeyableReplacesStable/v1" ) - register_memo_key_function( + register_memo_type( SameStableTypeId, stable_type_id="test.NotMemoKeyableReplacesStable/v1", ) @@ -1422,8 +1393,8 @@ class SameStableTypeId: _dummy_fn, (SameStableTypeId,), {}, [] ) finally: - unregister_memo_key_function(Entry) - unregister_memo_key_function(SameStableTypeId) + unregister_memo_type(Entry) + unregister_memo_type(SameStableTypeId) @@ -1444,7 +1415,7 @@ class SameStableTypeId: ), ], ) -def test_register_memo_key_function_rejects_invalid_forms( +def test_register_memo_type_rejects_invalid_forms( args: tuple[Any, ...], kwargs: dict[str, Any], match: str, @@ -1453,10 +1424,10 @@ class Entry: pass with pytest.raises(TypeError, match=match): - register_memo_key_function(Entry, *args, **kwargs) + register_memo_type(Entry, *args, **kwargs) -def test_register_memo_key_function_rejects_state_fn_without_key_function() -> None: +def test_register_memo_type_rejects_state_fn_without_key_function() -> None: class Entry: pass @@ -1466,10 +1437,58 @@ def state_fn(obj: Entry, prev_state: object) -> object: kwargs: Any = {"state_fn": state_fn} with pytest.raises(TypeError, match="state_fn requires a memo key function"): - register_memo_key_function(Entry, **kwargs) + 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) -def test_unregister_memo_key_function_clears_key_function_and_stable_type_id() -> None: + # 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 @@ -1479,12 +1498,12 @@ def __init__(self, value: object) -> None: self.value = value try: - register_memo_key_function( + register_memo_type( RegisteredOnly, lambda entry: ("registered", entry.value), stable_type_id="test.UnregisterCombined/v1", ) - register_memo_key_function( + register_memo_type( SameStableTypeId, lambda entry: ("registered", entry.value), stable_type_id="test.UnregisterCombined/v1", @@ -1496,15 +1515,15 @@ def __init__(self, value: object) -> None: fingerprint_call(_dummy_fn, (SameStableTypeId,), {}, []) ) - unregister_memo_key_function(RegisteredOnly) + 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_key_function(RegisteredOnly) - unregister_memo_key_function(SameStableTypeId) + unregister_memo_type(RegisteredOnly) + unregister_memo_type(SameStableTypeId) class OldEntry: def __init__(self, value: object) -> None: @@ -1521,18 +1540,18 @@ def __coco_memo_key__(self) -> object: return ("entry", self.value) try: - register_memo_key_function(OldEntry, stable_type_id="test.UnregisterStable/v1") - register_memo_key_function(NewEntry, stable_type_id="test.UnregisterStable/v1") + 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_key_function(OldEntry) + unregister_memo_type(OldEntry) assert fingerprint_call(_dummy_fn, (OldEntry(1),), {}, []) != fingerprint_call( _dummy_fn, (NewEntry(1),), {}, [] ) finally: - unregister_memo_key_function(OldEntry) - unregister_memo_key_function(NewEntry) + unregister_memo_type(OldEntry) + unregister_memo_type(NewEntry) def test_none_declared_stable_type_id_falls_back_to_registered_or_module_qualname() -> None: @@ -1551,7 +1570,7 @@ def __coco_memo_key__(self) -> object: stable_type_id = "test.NoneDeclarationFallback/v1" try: - register_memo_key_function(NoneId, stable_type_id=stable_type_id) + 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] == ( @@ -1560,27 +1579,29 @@ def __coco_memo_key__(self) -> object: None, ) finally: - unregister_memo_key_function(NoneId) + unregister_memo_type(NoneId) -def test_register_memo_key_function_validation_and_public_export() -> None: +def test_register_memo_type_validation_and_public_export() -> None: import cocoindex as coco - assert coco.register_memo_key_function is register_memo_key_function + 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 "register_memo_type_identifier" not in coco.__all__ assert not hasattr(coco, "register_memo_type_identifier") - assert "register_not_memo_keyable" not in coco.__all__ 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"): - register_memo_key_function( + coco.register_memo_type( cast(Any, object()), stable_type_id="test.Invalid/v1" ) - def test_cycles_are_supported_and_deterministic() -> None: # Self-cycle list a: Any = [] From 573a6940afd46e69f2b7b38613f9c1f77dfa69c8 Mon Sep 17 00:00:00 2001 From: Max Rong Date: Sun, 30 Aug 2026 00:40:49 +0000 Subject: [PATCH 16/16] feat(memo): expose register_memo_type in public API and enforce unregister boundaries - Re-export register_memo_type in top-level cocoindex namespace and __all__. - Update user-facing documentation and function docstrings to feature register_memo_type. - Keep unregister_memo_type and unregister_memo_key_function internal-only for test teardown. - Enforce public API boundaries with positive export and negative unregister assertions. --- .../docs/advanced_topics/memoization_keys.mdx | 31 +- .../docs/programming_guide/function.mdx | 2 +- python/cocoindex/_internal/api.py | 2 + python/cocoindex/_internal/function.py | 2 +- .../cocoindex/_internal/memo_fingerprint.py | 24 +- .../tests/internal/test_memo_fingerprint.py | 300 +++++++++++------- 6 files changed, 226 insertions(+), 135 deletions(-) diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index 6cda13056..7b32cb9c2 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -98,14 +98,14 @@ If you can't add `__coco_memo_key__` (stdlib / third-party types), register a ke ```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. @@ -142,13 +142,13 @@ class SourceEntry: # Moved to new_package.models 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_key_function()`: +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_key_function( +coco.register_memo_type( ExternalRecord, lambda record: (record.id, record.version), stable_type_id="com.example.ExternalRecord/v1", @@ -157,9 +157,11 @@ coco.register_memo_key_function( 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 propagates to subclasses. A subclass can override the ID by declaring its own. -- **Precedence**: A declared `__coco_memo_type_id__` (own or inherited) takes precedence over any registered stable type ID. -- **MRO registration**: Registered stable type IDs also apply to subclasses across the MRO: the most specific registered type wins. A subclass can override an inherited registered ID by declaring or registering its own. +- **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. @@ -167,14 +169,13 @@ Types using pickle fallback fingerprint raw bytes without type namespacing. Sett ### Class objects as inputs -A class object is the class value itself (such as `ProductRow`), not an instance (such as `ProductRow(...)`). +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. -When CocoIndex fingerprints a class object: +To customize memoization for class objects globally, register a memo key function on the class's custom metaclass or on `type`. -* It uses the class's stable type ID (declared with `__coco_memo_type_id__` or registered with `register_memo_key_function()`) when defined, or falls back to its module and qualified name (`module.QualName`). -* It does not call `__coco_memo_key__` or `__coco_memo_state__` methods defined on the class or its metaclass. -* It fingerprints only the class identity, not its schema or attributes. Include class details with `memo_key=` if the function depends on them. -* To customize memoization for class objects globally, register a memo key function on the class's custom metaclass or on `type` (which takes precedence over stable type IDs). +:::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=` @@ -293,7 +294,7 @@ 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 together with a callable `key_fn` to `register_memo_key_function`; `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: +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 @@ -310,7 +311,7 @@ def path_state( memo_valid = not coco.is_non_existence(prev_state) and new_state == prev_state return coco.MemoStateOutcome(state=new_state, memo_valid=memo_valid) -coco.register_memo_key_function( +coco.register_memo_type( Path, path_key, state_fn=path_state, diff --git a/docs/src/content/docs/programming_guide/function.mdx b/docs/src/content/docs/programming_guide/function.mdx index 7d3c8a4d8..b53a48cb4 100644 --- a/docs/src/content/docs/programming_guide/function.mdx +++ b/docs/src/content/docs/programming_guide/function.mdx @@ -185,7 +185,7 @@ Only `deps` is snapshotted at decoration time; regular function arguments are fi 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_key_function(..., stable_type_id=...)`, before any `@coco.fn(..., deps=...)` or `@coco.fn.as_async(..., deps=...)` decorator that references those values. +- 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. ::::: diff --git a/python/cocoindex/_internal/api.py b/python/cocoindex/_internal/api.py index fe9b6267a..5ec6a97c8 100644 --- a/python/cocoindex/_internal/api.py +++ b/python/cocoindex/_internal/api.py @@ -105,6 +105,7 @@ memo_fingerprint, prev_type_id, register_memo_key_function, + register_memo_type, NotMemoKeyable, ) @@ -946,6 +947,7 @@ class Cursor: "memo_fingerprint", "prev_type_id", "register_memo_key_function", + "register_memo_type", "NotMemoKeyable", # .pending_marker "MaybePendingS", diff --git a/python/cocoindex/_internal/function.py b/python/cocoindex/_internal/function.py index 35fc04114..d6e83da3f 100644 --- a/python/cocoindex/_internal/function.py +++ b/python/cocoindex/_internal/function.py @@ -2047,7 +2047,7 @@ def __call__( # type: ignore[misc] 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_key_function(..., stable_type_id=...)`` + 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. diff --git a/python/cocoindex/_internal/memo_fingerprint.py b/python/cocoindex/_internal/memo_fingerprint.py index e011eaa8d..89cf031c2 100644 --- a/python/cocoindex/_internal/memo_fingerprint.py +++ b/python/cocoindex/_internal/memo_fingerprint.py @@ -189,9 +189,10 @@ 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 @@ -207,6 +208,8 @@ def _type_identity_parts( 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) @@ -244,7 +247,7 @@ def _canonicalize_registered_memo_key( 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) + 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 @@ -292,7 +295,7 @@ def _canonicalize_class_object( *_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)), + ("seq", _type_identity_parts(cls, cls_registry, is_class_object=True)), ) @@ -386,6 +389,7 @@ def register_memo_type( typ: type, key_fn: None = None, *, + state_fn: None = None, stable_type_id: str, ) -> None: ... @@ -413,8 +417,7 @@ def register_memo_type( if not isinstance(typ, type): raise TypeError( - "register_memo_type() expects typ to be a type, " - f"got {type(typ).__name__}" + 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: @@ -422,13 +425,10 @@ def register_memo_type( "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" - ) + raise TypeError("register_memo_type() requires a key_fn or stable_type_id") elif not callable(key_fn): raise TypeError( - "register_memo_type() key_fn must be callable, " - f"got {type(key_fn).__name__}" + 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( @@ -460,7 +460,9 @@ def register_memo_key_function( """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) + register_memo_type( + typ, key_fn=key_fn, state_fn=state_fn, stable_type_id=stable_type_id + ) def register_not_memo_keyable(typ: type) -> None: diff --git a/python/tests/internal/test_memo_fingerprint.py b/python/tests/internal/test_memo_fingerprint.py index 5473a4124..b7efe53da 100644 --- a/python/tests/internal/test_memo_fingerprint.py +++ b/python/tests/internal/test_memo_fingerprint.py @@ -350,14 +350,6 @@ class OrdinaryStringSourceEntry: None, ) - class NonStringSourceEntry: - __coco_memo_type_id__ = 12345 - - assert _memo_fingerprint._type_identity_parts(NonStringSourceEntry, None) == ( - _memo_fingerprint.canonical_module_name(NonStringSourceEntry), - NonStringSourceEntry.__qualname__, - ) - def test_pydantic_stable_type_id_allows_renamed_model_reuse() -> None: try: @@ -429,15 +421,9 @@ 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" - ) + 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,), {}, [] @@ -789,15 +775,25 @@ def test_raw_class_object_ignores_registered_object_memo_key() -> None: class Entry: pass - expected = fingerprint_call(_dummy_fn, (Entry,), {}, []) + 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) - assert fingerprint_call(_dummy_fn, (Entry,), {}, []) == expected + 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) @@ -908,9 +904,7 @@ def base_key(entry: OldBase | MovedBase) -> object: unregister_memo_type(MovedBase) -def test_register_memo_type_registers_stable_type_id_without_key_function() -> ( - None -): +def test_register_memo_type_registers_stable_type_id_without_key_function() -> None: class OldEntry: def __init__(self, value: object) -> None: self.value = value @@ -961,15 +955,11 @@ def __coco_memo_key__(self) -> object: return ("entry", self.value) try: - register_memo_type( - Parent, stable_type_id="test.RegisteredParent/v1" - ) + 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" - ) + 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( @@ -991,7 +981,9 @@ def __coco_memo_key__(self) -> object: unregister_memo_type(SameStableTypeId) -def test_declared_stable_type_id_is_inherited_and_takes_precedence_over_registered() -> None: +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" @@ -1073,6 +1065,7 @@ class MovedSub(MovedBase): _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: @@ -1131,7 +1124,9 @@ class TwinSub(Base): unregister_memo_type(Base) -def test_subclass_registered_stable_type_id_overrides_registered_base_identity() -> None: +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 @@ -1169,89 +1164,117 @@ class TwinSub(Base): unregister_memo_type(Base) -def test_register_memo_type_full_replacement_lifecycle() -> None: +def test_register_memo_type_replaces_and_clears_omitted_state_fn() -> None: @dataclasses.dataclass class Entry: value: int ignored: str - @dataclasses.dataclass - class TwinV1Shook: - value: int - ignored: str - - @dataclasses.dataclass - class TwinV2Hook: - value: int - ignored: str - - @dataclasses.dataclass - class TwinV2Dataclass: - value: int - ignored: str - - id_v1 = "test.Lifecycle/v1" - id_v2 = "test.Lifecycle/v2" - 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( - TwinV1Shook, + Entry, lambda e: ("key", e.value), state_fn=state_fn, - stable_type_id=id_v1, - ) - register_memo_type( - TwinV2Hook, lambda e: ("key", e.value), stable_type_id=id_v2 + stable_type_id="test.Replace/v1", ) - register_memo_type(TwinV2Dataclass, stable_type_id=id_v2) + 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"),), {}, []) - # 1. Full registration (key_fn, state_fn, stable_id_v1): + # 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), - state_fn=state_fn, - stable_type_id=id_v1, + stable_type_id="test.Replace/v1", ) - methods: list[Any] = [] - fp1 = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) - assert fp1 == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) # key_fn active (ignores 'ignored') - assert fp1 == fingerprint_call(_dummy_fn, (TwinV1Shook(1, "a"),), {}, []) # identity is id_v1 and tag is shook - assert len(methods) == 1 # state_fn active + 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 - # 2. Replace with (key_fn, stable_id_v2): clears state_fn, updates stable_id to id_v2, keeps key_fn + # 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=id_v2, + stable_type_id="test.Replace/v1_updated", ) - methods = [] - fp2 = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, methods) - assert len(methods) == 0 # state_fn cleared (tag is now hook) - assert fp2 != fp1 # stable_id changed and state_fn cleared - assert fp2 == fingerprint_call(_dummy_fn, (TwinV2Hook(1, "a"),), {}, []) # matches TwinV2Hook - assert fp2 == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) # key_fn still active - - # 3. Replace with stable_id_v2 only: clears key_fn while keeping stable_id_v2 - register_memo_type(Entry, stable_type_id=id_v2) - fp3 = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) - assert fp3 != fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) # key_fn cleared (dataclass default participates) - assert fp3 == fingerprint_call(_dummy_fn, (TwinV2Dataclass(1, "a"),), {}, []) # stable_id_v2 preserved - - # 4. Replace with key_fn only: clears stable_id (identity falls back to module.qualname) + 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)) - fp4 = fingerprint_call(_dummy_fn, (Entry(1, "a"),), {}, []) - assert fp4 == fingerprint_call(_dummy_fn, (Entry(1, "b"),), {}, []) # key_fn active again - assert fp4 != fingerprint_call(_dummy_fn, (TwinV1Shook(1, "a"),), {}, []) # stable_id cleared - assert fp4 != fingerprint_call(_dummy_fn, (TwinV2Hook(1, "a"),), {}, []) - assert fp4 != fingerprint_call(_dummy_fn, (TwinV2Dataclass(1, "a"),), {}, []) + 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(TwinV1Shook) - unregister_memo_type(TwinV2Hook) - unregister_memo_type(TwinV2Dataclass) + unregister_memo_type(SameIdEntry) + def test_intrinsic_hooks_and_declared_stable_id_beat_registration() -> None: declared_type_id = "test.DeclaredIntrinsic/v1" @@ -1368,7 +1391,6 @@ def state_fn(entry: Any, prev_state: object) -> MemoStateOutcome: unregister_memo_type(NewEntry) - def test_register_not_memo_keyable_replaces_stable_type_id_for_class_objects() -> None: class Entry: pass @@ -1377,9 +1399,7 @@ class SameStableTypeId: pass try: - register_memo_type( - Entry, stable_type_id="test.NotMemoKeyableReplacesStable/v1" - ) + register_memo_type(Entry, stable_type_id="test.NotMemoKeyableReplacesStable/v1") register_memo_type( SameStableTypeId, stable_type_id="test.NotMemoKeyableReplacesStable/v1", @@ -1397,7 +1417,6 @@ class SameStableTypeId: unregister_memo_type(SameStableTypeId) - @pytest.mark.parametrize( ("args", "kwargs", "match"), [ @@ -1440,7 +1459,9 @@ def state_fn(obj: Entry, prev_state: object) -> object: register_memo_type(Entry, **kwargs) -def test_register_and_unregister_memo_key_function_shortcuts_preserve_stable_type_id() -> None: +def test_register_and_unregister_memo_key_function_shortcuts_preserve_stable_type_id() -> ( + None +): @dataclasses.dataclass class Entry: value: int @@ -1458,36 +1479,46 @@ class Twin: # 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( + 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 + 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( + 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( + 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( + 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 + 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: @@ -1554,12 +1585,15 @@ def __coco_memo_key__(self) -> object: unregister_memo_type(NewEntry) -def test_none_declared_stable_type_id_falls_back_to_registered_or_module_qualname() -> None: +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] == ( @@ -1582,11 +1616,48 @@ def __coco_memo_key__(self) -> object: 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.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__ @@ -1598,9 +1669,24 @@ def test_register_memo_type_validation_and_public_export() -> None: 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( - cast(Any, object()), stable_type_id="test.Invalid/v1" + 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