Skip to content

feat(memo): add stable memo type identifiers - #2256

Open
MaxRong wants to merge 16 commits into
cocoindex-io:mainfrom
MaxRong:feat/custom-object-type-identifier
Open

feat(memo): add stable memo type identifiers#2256
MaxRong wants to merge 16 commits into
cocoindex-io:mainfrom
MaxRong:feat/custom-object-type-identifier

Conversation

@MaxRong

@MaxRong MaxRong commented Jul 5, 2026

Copy link
Copy Markdown
Contributor

Closes: #2229

Summary

This PR adds opt-in stable type IDs for memo fingerprints, allowing a semantic Python type to keep the same memo namespace across a module move or class rename.

By default, CocoIndex keeps the existing behavior: type-aware fingerprints include the Python type's module and qualified name.

Public API

Types controlled by the application can define a stable ID directly:

from dataclasses import dataclass

@dataclass
class SourceEntry:
    __coco_memo_type_id__ = "com.example.SourceEntry/v1"

    name: str
    version: int

Types that cannot be modified use the existing registration API:

from pathlib import Path

import cocoindex as coco

coco.register_memo_key_function(
    Path,
    lambda path: str(path.resolve()),
    stable_type_id="stdlib.pathlib.Path/v1",
)

A stable type ID can also be registered without a key function when CocoIndex already knows how to fingerprint the exact type:

coco.register_memo_key_function(
    SourceEntry,
    stable_type_id="com.example.SourceEntry/v1",
)

register_memo_key_function(...) is the single public registration API. It supports:

  • key function only;
  • key plus state function;
  • stable type ID only;
  • key/state functions plus stable type ID.

Each call replaces the exact type's complete registered key/state/stable-ID configuration. The previously proposed register_memo_type_identifier(...) API is not exposed.

Class objects

A class object means the class value itself—for example, ProductRow—rather than an instance such as ProductRow(...).

Class values honor explicit key/state registrations on their custom metaclass or type. Lookup stops before object. If there is no matching registration, CocoIndex uses the class's stable type ID or its module and qualified name.

Migration

Stable type IDs do not rewrite existing memo entries.

To preserve reuse across a move or rename:

  1. Add the stable type ID before the refactor.
  2. Execute affected memoized calls to seed the new namespace.
  3. Move or rename the class while retaining the same stable ID.

A cached parent does not execute nested memoized calls, so affected parents may need to be invalidated or reprocessed during seeding.

Breaking changes

Raw class values no longer call __coco_memo_key__() or __coco_memo_state__() attributes found on the class or its metaclass, including callable @classmethod or @staticmethod forms.

Applications intentionally relying on that behavior should migrate to:

  • per-function memo_key= for local behavior;
  • registration on a dedicated custom metaclass for one class family;
  • registration on type for intentional process-wide class behavior.

Existing instance memo-key and memo-state methods are unaffected.

Scope / non-goals

This PR intentionally does not:

  • migrate memo entries created before a stable type ID was added;
  • make arbitrary pickle-fallback objects refactor-stable without __coco_memo_key__ or a registered memo-key function;
  • fingerprint class schema or mutable class-level state automatically;
  • allow a registration on object to customize raw class values;
  • change function logic identity, which still uses the function's module and qualified name.

Testing

  • uv run pytest python/`
    skipped
  • uv run mypy
  • cd docs && npm ci && npm run build

@MaxRong
MaxRong marked this pull request as draft July 5, 2026 22:32
@MaxRong
MaxRong force-pushed the feat/custom-object-type-identifier branch from 1df93b5 to 9b17792 Compare July 5, 2026 23:35
@MaxRong
MaxRong marked this pull request as ready for review July 5, 2026 23:44

@georgeh0 georgeh0 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for making this PR!


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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I failed to understand this paragraph.

@MaxRong MaxRong Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I should make it more clear.

  • Stable type IDs only help for memo keys created after the ID exists. They don’t rewrite old memo cache entries. Additionally, If Cocoindex falls back to pickle for a plain object with no memo key hook the stable ID is not used. So if a refactor-stable identity is needed, the user should define one with __coco_memo_key__ or register_memo_key_function(...) unless key function includes namespace.
  • And, for class objects, if explicit coco_memo_key is defined that hook overrides default behavior, so it's important for users to include the stable type id in the return.

@georgeh0 georgeh0 Jul 7, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was asking for the exact last paragraph. I think I can understand the two paragraphs above.

I still don't really understand the last paragraph, until I read the implementation: I just know what does class-object really mean here (I thought it simply meant an object constructed by a class type).

Observations and suggestions:

  • When someone says they fails to understand a part of the document, consider if the document needs to be improved.
  • Usually using examples will make a document easier to understand.
  • This paragraph uses the term "hook" multiple times, but it's never defined or used anywhere else in the doc. So readers easily get lost.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

noted, I'll make this a lot more clear, thank you for the detailed tips!

Comment on lines +107 to +108
coco.register_memo_key_function(Path, lambda path: str(path))
coco.register_memo_type_identifier(Path, "stdlib.pathlib.Path/v1")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An alternative way is to allow to register both in a single function (e.g. users can provide both identifier and key function, and can also only provide one of them).

What do you think about pros and cons of both approaches?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I kept type identity and key function separate because it felt more explicit, especially since you can reasonably want one without the other. That said, I agree a unified function seems very reasonable.

If you think one entry point would be more ergonomic, I’m happy to add it. I’d probably keep the explicit APIs available too, and implement a register_memo_type function as a convenience wrapper over them.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it felt more explicit, especially since you can reasonably want one without the other

I think providing one without the other is always possible even under the same API

I prefer a single unified API since they're for the same purpose. Separated APIs + convenient wrapper will broaden the API surface unnecessarily, resulting in choice paralysis and diverged coding styles (e.g. consider if we want to search for registrations for given types, we need to remember searching for both APIs).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

understood, I completely agree. Will work in the direction of extending the register_memo_key function in that case.


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.

@georgeh0 georgeh0 Jul 7, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was asking for the exact last paragraph. I think I can understand the two paragraphs above.

I still don't really understand the last paragraph, until I read the implementation: I just know what does class-object really mean here (I thought it simply meant an object constructed by a class type).

Observations and suggestions:

  • When someone says they fails to understand a part of the document, consider if the document needs to be improved.
  • Usually using examples will make a document easier to understand.
  • This paragraph uses the term "hook" multiple times, but it's never defined or used anywhere else in the doc. So readers easily get lost.

"""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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

curious why not directly return ("__coco_memo_type_id__", identifier) here

@MaxRong MaxRong Jul 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wanted to tag the stable ID case explicitly to distinguish from the normal (module, qualname). There's also some minor collision risk but a non-issue for average users likely, can simplify if preferred for maintainability

if raw is None:
continue
hook = getattr(cls, _MEMO_KEY_ATTR, None)
if isinstance(raw, (classmethod, staticmethod)):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I understand the motivation, but feel it's tricky if we expect users to provide a @staticmethod/@classmethod for __coco_memo_key__ for class-level memo key extraction: when come to the memo key extraction for an instance of the class, it'll end up with the same value for all instances of the class. It's a footgun.

For the class itself, maybe we should never call __coco_memo_key__.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

great point, I overlooked this and tried to get too clever. Will definitely simplify and take your suggestion.

Comment on lines +107 to +108
coco.register_memo_key_function(Path, lambda path: str(path))
coco.register_memo_type_identifier(Path, "stdlib.pathlib.Path/v1")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it felt more explicit, especially since you can reasonably want one without the other

I think providing one without the other is always possible even under the same API

I prefer a single unified API since they're for the same purpose. Separated APIs + convenient wrapper will broaden the API surface unnecessarily, resulting in choice paralysis and diverged coding styles (e.g. consider if we want to search for registrations for given types, we need to remember searching for both APIs).

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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what does this mean? is it defined below?


- **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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i see multiple definitions for similar stuff
stable semantic type identifiers,
Stable type identifiers
semantic type ID

are they the same thing?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I will standardize to 'stable type IDs' and make documentation a lot more clear.

@MaxRong

MaxRong commented Jul 10, 2026

Copy link
Copy Markdown
Contributor Author

Thanks again for the detailed review. I pushed the follow-up commit:

a2d0a2ae feat(memo): implement reviewer feedback, fix bugs

and also edited the PR body.

What changed

  • Unified registration: register_memo_key_function(...) now handles
    key, state, stable-ID-only, and combined registration. Each call replaces
    the exact type's complete registered configuration. The separate
    register_memo_type_identifier(...) API was removed.
  • Class-object behavior: class values such as ProductRow no longer
    invoke memo-key/state attributes found on the class itself. Explicit
    custom-metaclass and type registrations still work, while lookup stops
    before object.
  • Canonicalization: custom memo-key fragments now share their parent's
    traversal, preserving parent cycles and aliases. Traversal-local strong
    references prevent temporary object-ID reuse from creating false references.
  • Identity handling: registration is exact by Python type-object
    identity, and the registries do not keep registered types alive. Equal or
    unhashable metaclasses are covered.
  • Timing and state: the implementation and docs now distinguish
    argument call time, context builder.provide() time, and deps decoration
    time. Updated memo state is persisted when a cached result remains reusable.
  • Documentation: the guide now defines class objects explicitly,
    removes undefined “hook” terminology, explains stable-ID seeding and
    replacement, and adds concrete key/state examples.
  • Regression coverage: added focused tests for class dispatch,
    object exclusion, registration replacement, cycles, aliases, temporary
    fragments, hostile equality/hash behavior, state timing, and rename reuse.

Compatibility note

Raw class values no longer call class-level __coco_memo_key__() /
__coco_memo_state__() methods, including @classmethod or @staticmethod
forms. Instance methods are unaffected.

Raw-class customization should use per-function memo_key=, a dedicated
metaclass registration, or intentional process-wide registration on type.

Verification

  • 6 focused class-dispatch/registration regressions passed;
  • 2 focused object-boundary/temporary-fragment regressions passed;
  • documentation build passed with 95 pages;
  • documentation checks passed.

I would especially appreciate another look at:

  • registration replacement(behaving as expected)?
  • revised documentation(migration guidance, clarity).

Comment on lines +15 to +22
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).

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice summarization!

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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel this is not a major case, and placing it as a first item (with many words) will make it significantly harder to understand.

I think we can remove it. We only need to add one sentence to "2" (which will become "1") to clarify, e.g. something like:

1. If the object implements **`__coco_memo_key__()`**, CocoIndex uses its return value. This applies to class instances, but not class objects (`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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

parents—or

It's really hard to distinguish em dash and a regular dash (especially under monospace fonts). parents—or looks like one hyphenated word. Please either avoid em dash, or add spaces before/after the em dash.

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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel we can make this section simpler.

The "seeding execution" described in this section seems not really necessary for most cases. Once users start to set or change __coco_memo_type_id__, previous memoizations are invalidated any way. So IMO "seeding execution" is not really meaningful.

We have two options here:

  • We can suggest users that, if memo stability is really important and you cannot rule out potential refactor in the future, set a __coco_memo_type_id__. That's the only way to avoid a memo invalidation in future migration.

  • We can also allow users directly carry over the auto-generated type ID we used before when __coco_memo_type_id__ was not set, e.g. if foo.Class1 is moved to a another module, users can right something like:

    __coco_memo_type_id__ = coco.prev_type_id('foo', 'Class1')

    This is more friendly for migration.

What do you think?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review!

I agree with what you've said, your suggestions would make this process a lot clearer. Let me clear up the comments and implement the prev_type_id function.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good. A prev_type_id function is also the preferred approach in my mind. Please implement it. Thanks!

@MaxRong
MaxRong force-pushed the feat/custom-object-type-identifier branch from ee7c869 to a2d0a2a Compare July 16, 2026 02:30
validation (see ``_canonicalize``).
def register_memo_key_function(
typ: type,
key_fn: _KeyFn | object = _KEY_FN_UNSET,

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not simply use None as unset?

Comment on lines +213 to +215
if "__coco_memo_type_id__" in typ.__dict__:
return _validate_stable_type_id(
typ.__dict__["__coco_memo_type_id__"],

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"__coco_memo_type_id__" in typ.__dict__ and typ.__dict__["__coco_memo_type_id__"] are doing redundant field loop.

I think we can simply do getattr(typ, '__coco_memo_type_id__', None) once and check the result.


_memo_fns: dict[type, _MemoFns] = {}
_memo_fns: dict[int, tuple[weakref.ReferenceType[type], _MemoFns]] = {}
_stable_type_ids: dict[int, tuple[weakref.ReferenceType[type], str]] = {}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think a cleaner way is to simply add a field to _MemoFns for the new type ID. We only need a single dict. These 3 fields are usually needed set and used together. The logic will be more clear since all registered information for 3rd party types are kept together.

We only need to give _MemoFns a broader name, e.g. _MemoTypeRegistry.

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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good. A prev_type_id function is also the preferred approach in my mind. Please implement it. Thanks!

## 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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a nice elaboration for the existing short paragraph.

This clearly answers a question of "when is fingerprinting needed" (rather than "how"), and the first paragraph of this doc aims at covering it. Maybe can elaborate in the 1st paragraph there (e.g. split into separate bullets to be clearer), and keep this one concise (since it's for "how").

@MaxRong

MaxRong commented Jul 18, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the round of feedback! I addressed the concerns:

  • Replaced the custom omitted-key_fn sentinel with None; explicit None now supports stable-ID-only registration.
  • Consolidated coco_memo_type_id resolution into one typ.dict.get(..., _MISSING) lookup. I used the exact class dictionary rather than getattr to avoid unintentionally inheriting a parent class’s ID and to distinguish a missing attribute from an invalid explicit None.
  • Unified key, state, and stable-ID registrations into one _MemoTypeRegistry record and one identity-keyed weak-reference table.
  • Added prev_type_id(module, qualname) for types that have already been memoized under their automatic (module, qualname) identity and are later moved or renamed. Assigning its result to coco_memo_type_id, or passing it as stable_type_id, keeps emitting the previous automatic identity so compatible existing memo entries remain addressable without copying or migrating records. The documentation covers canonical module/qualified-name values, nested classes, main, limitations, and collision risks.
  • Moved fingerprint timing details into the guide introduction


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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the case when the type is in the registry, we need the other fields in the registry too. Ideally we only need to lookup the registry once for the type, in _canonicalize().

BTW, when a type both __coco__* members also registered in the registry, the __coco__* members should take precedence.

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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can directly use None instead of _STABLE_TYPE_ID_MISSING. None isn't a valid value for __coco_memo_type_id__ any way.

Comment on lines +161 to +166
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:]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is on the hot path (_canonicalize() is called very frequently). We have better ways than computing the parts out of the str on the fly. Can we simply store the module and qualname as new field slots?

Comment on lines +168 to +174
@property
def module(self) -> str:
return self._identity_parts()[0]

@property
def qualname(self) -> str:
return self._identity_parts()[1]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems they're not used anywhere.

@MaxRong
MaxRong force-pushed the feat/custom-object-type-identifier branch from 9ae2849 to d0f0880 Compare July 25, 2026 17:42
@MaxRong

MaxRong commented Jul 25, 2026

Copy link
Copy Markdown
Contributor Author

Just rebased everything, newest changes on the last commit. Thanks for the new round of feedback! Here is the changes:

  • The exact runtime type's registry record is now resolved once in
    _canonicalize() and passed through the identity/canonicalization
    path instead of being looked up repeatedly.
    • __coco_memo_key__ / __coco_memo_state__ members
      take precedence over registered key/state functions.
    • A non-None class-declared __coco_memo_type_id__ now takes
      precedence over the registry's stable type ID, including when an
      MRO-selected registered key function handles the value.
    • Removed the separate missing-value sentinel
    • _PreviousTypeId now stores (module, qualname) directly, so
      the hot path no longer reparses its string payload. The unused
      public-looking accessors were removed, while copy/deepcopy/pickle
      round trips remain covered.
    • I kept a single typ.__dict__.get("__coco_memo_type_id__")
      lookup rather than getattr(...) because stable IDs are
      intentionally exact-type-only and should not be inherited from a
      parent class.
    • Updated the guide to reflect the corrected class-declared-ID
      precedence and added focused regressions for precedence and registry
      lookup counts.

Comment on lines +166 to +182
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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel we don't have to have these two separate util functions just to validate arguments module and qualname for prev_type_id(). module and qualname are already annotated with type hint str in prev_type_id() and we usually just trust that (otherwise too many things to validate everywhere).

Also IMO we don't have to check against "" too: there can be many different types of wrong values that don't match the actual previous type, and "" isn't special from others. The contract of our API is "if the passed-in names match previous one, we keep memoization key stability", but we have no obligation to validate they actually match (since it's impossible).


def __setattr__(self, name: str, value: object) -> typing.NoReturn:
raise AttributeError(f"{type(self).__name__} is immutable")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since __setattr__ is blocked, we may also want to consider blocking __delattr__ to make it really immutable.

Comment on lines +271 to +283
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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_canonicalize_key_fragment() is a trivial wrapper around _canonicalize(). Seems this layer provides no value. We may eliminate it to simplify.

Comment on lines +257 to +260
stable_type_id = _validate_stable_type_id(
stable_type_id,
source=f"{_memo_type_label(typ)}.__coco_memo_type_id__",
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IMO this is another unnecessary validation and it's on hot path. We can't really meaningfully validate that, so we don't have to.

Comment on lines +114 to +115
### Use stable type IDs across refactors

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This entire section has 150 lines now, which is too verbose. I think the underlying semantics is deadly simple: if you move/rename your type, assign a stable type ID with the previous name. We don't need a section with 150 lines.

…dation

- 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.
…dence

- 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.
- 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.
@MaxRong

MaxRong commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

Hello, I update the PR based on feedback:

  • Simplified prev_type_id() and _PreviousTypeId (added __delattr__, put
    _identity_parts in slots, dropped the extra validation helpers and wrapper).
  • Pruned memoization_keys.mdx down to ~45 lines focused on the main migration patterns, and
    clarified metaclass/type precedence.
  • Swapped the weakref registry table for a normal dict[type, _MemoTypeRegistry] and removed
    the old weak-identity tests.

stable_type_id = registry.stable_type_id

if isinstance(stable_type_id, _PreviousTypeId):
return stable_type_id._identity_parts

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually here doing another simple isinstance(stable_type_id, str) check is warranted, since we need type check any way, and there's to make the Python type checker validate the type for __coco_memo_type_id__.

"register_memo_key_function() state_fn must be callable, "
f"got {type(state_fn).__name__}"
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since we verified other arguments here, let's also verify type of stable_type_id for consistency.

tagged first slot keeps stable type IDs disjoint from ordinary module
names; ``None`` fills the qualname slot.
"""
stable_type_id = typ.__dict__.get("__coco_memo_type_id__")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Following up on our earlier exchange (I originally suggested getattr just for brevity; you kept __dict__ to avoid unintentional inheritance): after more thought, let's switch to getattr, i.e. stable type IDs are inherited by subclasses.

Reasons:

  • There are valid use cases for inheritance: a family of classes sharing one base, where the base fully decides how memo keys are computed and subclasses are implementation details (e.g. the same data represented in different storages). Since __coco_memo_key__ is inherited but the type ID is not, such a family today silently gets fragmented namespaces unless every subclass redeclares the ID.
  • It gives one uniform dispatch rule for all __coco_* members: normal attribute lookup on the object first (own or inherited — we can't distinguish anyway), registered configuration as fallback, then module + qualname. Today methods inherit but the field doesn't; that asymmetry is hard to teach.
  • It matches ordinary Python semantics for a class attribute; the __dict__-only lookup is the surprising behavior.
  • Docs get simpler: the "exact-type-only, subclasses don't inherit" caveat goes away.

On the aliasing concern that motivated __dict__ (parent/child instances sharing one namespace): declaring an ID on a base class is a deliberate statement that the family shares a namespace. A subclass that diverges semantically should override the ID: same discipline as overriding __eq__ when equality semantics change. For dataclass/Pydantic families the field tuples keep payloads distinct anyway. So this is acceptable opt-in behavior; just document it.

Two things to handle together with the switch:

  1. prev_type_id on a non-final base class: each subclass's automatic identity was its own (module, qualname), so a subclass inheriting the parent's marker would emit the parent's old identity and silently mismatch the subclass's own old entries. Add one sentence to the docs: declare prev_type_id on each moved class; like any class attribute, a base-class declaration propagates to subclasses.
  2. Precedence: pin down that a declared __coco_memo_type_id__ (own or inherited) takes precedence over a registered stable_type_id, including the subclass's own exact-type registration: keep the uniform "declared beats registered" rule. Please state it in the docs and cover it with a small test.

Let's make stable-ID-only registrations MRO-aware as well, so declared and registered stable type IDs both apply to subclasses (most specific registered type wins; a subclass overrides by declaring or registering its own). Today a registered ID covers subclasses only when the registration also carries a key function, which is a conditional scope rule that's hard to teach, and registration is the only mechanism available for third-party hierarchies, where the shared-family case matters most. Precedence stays mechanism-major, mirroring key extraction: declared ID (normal attribute lookup) → registered ID (MRO) → module + qualname.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, I'll implement this immediately.


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`).

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not precise. register_memo_key_function also affects this.

Comment on lines +370 to +371
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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since register_memo_key_function is now the way to register something that isn't a key function (and clears one), consider renaming to register_memo_type. We can keep register_memo_key_function taking the original signature as a shortcut / for backward compatibility.


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)."""

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: we can remove the (best-effort) here.

:::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.
:::

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Somewhere we may want to briefly mention the "full replace" semantics, e.g.

Each call replaces the type's entire registration: register the key function, state function, and stable type ID together in one call.

- 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.
…istration 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.
…mo_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.
…ister 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.
@MaxRong

MaxRong commented Aug 31, 2026

Copy link
Copy Markdown
Contributor Author

PR Update

Just pushed the latest patch of changes addressing the comments:

What Changed

  • Inherited declared stable type IDs via getattr with MRO registration fallback (Declared -> Registered MRO -> module.QualName).
  • Renamed the primary registration function to register_memo_type with full-replacement semantics for key function, state function, and stable type ID.
  • Retained register_memo_key_function as a backward-compatible shortcut taking its original signature.
  • Enforced string type checks for stable type IDs, prevented raw class MRO traversal from reaching object, and updated docs and tests.

Key Decisions

  • When a base class registers a key function, subclasses inherit and execute that key function, but any declared or registered stable type ID on the subclass overrides the identity namespace. If the subclass defines no stable ID, it falls back to the base owner's identity.
  • Considered consolidating registered key-owner and stable-ID resolution into a single MRO pass, but decided to keep _type_identity_parts() decoupled. Key extraction and type-identity namespacing are distinct responsibilities (consumed independently by intrinsic hooks, dataclasses, and Pydantic models), and thought combining them would unnecessarily couple the dispatch loop especially since declared IDs already resolve in O(1) via getattr and bypass MRO traversal entirely. Additionally, python MROs are rarely deep.
  • For the legacy shortcut register_memo_key_function, replaces key_fn/state_fn but intentionally preserves any existing stable_type_id. The legacy signature cannot accept stable_type_id, so clearing it on key-function registration would be unexpected and detrimental for those using the legacy function.
  • Maintained symmetrical behavior in the internal test fixture functions unregister_memo_type and unregister_memo_key_function. Kept intentionally private to reduce public API surface (and discourage improper patterns).

CI Note

The standard Windows Python 3.11 job timed out at the 30.0s mark on an untouched test (test_bound_method_memo_with_use_mount) due to runner I/O latency. Could we rerun the Github Actions suite?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FEATURE] customized type identifier to replace full-qualified module path + type name

3 participants