gemini-code-assist[bot] commented on code in PR #663: URL: https://github.com/apache/tvm-ffi/pull/663#discussion_r3576031961
########## python/tvm_ffi/dataclasses/_resolve_fields.py: ########## @@ -0,0 +1,462 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Forward-annotation resolution for ``@py_class`` registration. + +``@py_class`` registers a Python type before it resolves field annotations, so +self-referential and mutually recursive annotations can refer to classes that +are already present in the FFI type registry. This module owns that early +registration step, the per-module local namespace used by +``typing.get_type_hints``, and the queue of classes whose annotations must be +retried after later definitions become available. + +When a class is ready to finalize, this module lazily imports +``py_class.on_fields_resolved``. Keeping the finalizer import local preserves +the boundary between annotation resolution and field registration while avoiding +a top-level import cycle. +""" + +from __future__ import annotations + +import sys +import typing +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from .. import core +from .field import KW_ONLY + +if TYPE_CHECKING: + from ..core import TypeInfo +else: + TypeInfo = Any + +ResolvedFields = tuple[list[type], dict[type, dict[str, Any]]] + + +# --------------------------------------------------------------------------- +# Module-level state +# --------------------------------------------------------------------------- +# +# ``@py_class`` registration happens in two phases: +# +# Phase 1 (register_type_without_fields) +# Allocates a C-level type index and inserts the class into the +# global type registry. This must happen early so that self- +# referential and mutually-referential annotations can resolve +# the class via ``TypeSchema.from_annotation()``. Phase 1 always +# succeeds or raises immediately for invalid parents. Decorator +# options are stored by py_class.py on ``TypeInfo._decorator_args``. +# +# Phase 2 (py_class.on_fields_resolved) +# Resolves string annotations via ``typing.get_type_hints``, +# returns per-owner resolved hints, and asks py_class.py to +# materialize ``Field`` objects, register fields/methods/type +# attributes with the Cython layer, and install dataclass dunders. +# +# If ``get_type_hints`` raises ``NameError`` (forward reference +# not yet defined), the class is added to ``_PENDING_CLASSES`` +# and retried after each successful phase-2. If phase-2 fails +# for any other reason, ``rollback_registration`` undoes phase-1 +# so the type key can be reused. +# --------------------------------------------------------------------------- + + +@dataclass +class _PendingClass: + """Resolution state for a class waiting on unresolved annotations. + + Only data needed to retry annotation resolution is stored here. Decorator + options live on ``type_info._decorator_args`` so deferred and immediate + finalization use the same source of truth. + """ + + cls: type + type_info: TypeInfo + globalns: dict[str, Any] + + +#: Classes whose phase-2 (field registration) was deferred because +#: ``typing.get_type_hints`` raised ``NameError`` on an unresolved +#: forward reference. Retried after each successful phase-2 via +#: :func:`flush_pending`. +_PENDING_CLASSES: list[_PendingClass] = [] + +#: Per-module mapping of ``class.__name__ → class`` for every +#: ``@py_class``-decorated type. Used as *localns* when resolving +#: annotations so that mutual references between classes in the same +#: module work even before the second class is assigned to the module +#: variable by Python. +_PY_CLASS_BY_MODULE: dict[str, dict[str, type]] = {} + + +class _KWOnlyAnnotation: + """Type-valued stand-in accepted by ``typing.get_type_hints``.""" + + +def _is_kw_only_annotation(annotation: Any) -> bool: + if annotation is KW_ONLY or annotation is _KWOnlyAnnotation: + return True + if isinstance(annotation, str): + return annotation == "KW_ONLY" or annotation.endswith(".KW_ONLY") + forward_arg = getattr(annotation, "__forward_arg__", None) + if isinstance(forward_arg, str): + return forward_arg == "KW_ONLY" or forward_arg.endswith(".KW_ONLY") + return False + + +# --------------------------------------------------------------------------- +# Phase 1: type registration +# --------------------------------------------------------------------------- + + +def _registered_type_info(cls: type) -> TypeInfo | None: + """Return the TypeInfo registered directly for *cls*, not inherited metadata.""" + info = core._type_cls_to_type_info(cls) + if info is not None: + return info + return cls.__dict__.get("__tvm_ffi_type_info__", None) + + +def register_type_without_fields(cls: type, type_key: str | None) -> TypeInfo: + """Register ``cls`` in the FFI type registry before resolving fields. + + The returned :class:`~tvm_ffi.core.TypeInfo` has a type index, type key, and + Python class binding, but no field metadata yet. This early registration is + what lets ``TypeSchema.from_annotation`` resolve references to ``cls`` while + the decorator is still running. + + The class is also inserted into the module-local annotation namespace used + by :func:`resolve_type_hints_by_owner`, allowing sibling ``@py_class`` + declarations in the same module to reference each other before the second + class has been assigned to the module global by Python. + """ + parent_info = next( + (info for base in cls.__mro__[1:] if (info := _registered_type_info(base)) is not None), + None, + ) + if parent_info is None: + raise TypeError( + f"{cls.__name__} must inherit from a registered FFI Object type (e.g. tvm_ffi.Object)" + ) + if type_key is None: + type_key = f"{cls.__module__}.{cls.__qualname__}" + info = core._register_py_class(parent_info, type_key, cls) + setattr(cls, "__tvm_ffi_type_info__", info) + # Register in resolution namespace so sibling classes can find us. + _PY_CLASS_BY_MODULE.setdefault(cls.__module__, {})[cls.__name__] = cls + return info + + +def rollback_registration(cls: type, type_info: TypeInfo) -> None: + """Undo Python-side state from :func:`register_type_without_fields`. + + The C-level type index is permanently consumed (cannot be reclaimed), + but the Python-level registry dicts are cleaned up so a retry with + the same type key does not hit "already registered". The dataclass marker + is also removed because failed phase-2 finalization means the class is not a + usable FFI dataclass. + """ + # Remove from the Cython-level registry dicts (TYPE_KEY_TO_INFO, + # TYPE_CLS_TO_INFO, TYPE_INDEX_TO_INFO, TYPE_INDEX_TO_CLS). + core._rollback_py_class(type_info) # ty: ignore[unresolved-attribute] + # Remove from our own module-level resolution namespace. + _PY_CLASS_BY_MODULE.get(cls.__module__, {}).pop(cls.__name__, None) + for attr in ("__tvm_ffi_type_info__", "__ffi_is_dataclass__"): + try: + delattr(cls, attr) + except AttributeError: + pass + + +# --------------------------------------------------------------------------- +# Annotation resolution +# --------------------------------------------------------------------------- + + +def own_annotations(cls: type) -> dict[str, Any]: + """Return annotations declared directly on ``cls`` without MRO merging.""" + # Python 3.14+ (PEP 749): annotations are lazily evaluated via + # __annotate__ and no longer stored directly in __dict__. getattr() + # triggers evaluation and returns per-class annotations correctly. + # On Python < 3.14, getattr() follows MRO and returns *parent* + # annotations when the child has none — use __dict__ to avoid that. + if sys.version_info >= (3, 14): + return getattr(cls, "__annotations__", {}) + return cls.__dict__.get("__annotations__", {}) + + +def _field_owner_classes(cls: type) -> list[type]: + """Return local MRO entries whose annotations become fields on ``cls``. + + A Python subclass of a registered FFI parent may include unregistered mixin + bases before the registered parent. Those mixin annotations are owned by + the new Python type and should be registered together with ``cls``. Fields + already represented by the nearest registered parent are excluded. + """ + registered_parent = next( + (b for b in cls.__mro__[1:] if _registered_type_info(b) is not None), object + ) + represented = set(registered_parent.__mro__) + return [ + b + for b in reversed(cls.__mro__) + if b is not object and b not in represented and own_annotations(b) + ] + + +def _build_localns(cls: type, *, cross_module: bool = False) -> dict[str, Any]: + """Build the localns dict for resolving ``cls``'s annotations. + + By default, includes only classes from ``cls.__module__``, preserving + standard Python name resolution semantics. When ``cross_module=True``, + also includes classes from all other registered modules as a fallback + — this is needed when ``cls`` has a forward reference to a class in + another module that can't appear in ``cls.__module__``'s globals due + to a circular import (e.g. the target is imported only under + ``if TYPE_CHECKING:``). + + Cross-module entries are added with ``setdefault`` so same-module + classes and the class itself always take precedence over foreign + classes with the same ``__name__``. + """ + localns = dict(_PY_CLASS_BY_MODULE.get(cls.__module__, {})) + localns[cls.__name__] = cls + if cross_module: + for mod_name, mod_classes in list(_PY_CLASS_BY_MODULE.items()): + if mod_name == cls.__module__: + continue + for name, klass in mod_classes.items(): + localns.setdefault(name, klass) + return localns + + +def _resolve_own_type_hints( + owner: type, + globalns: dict[str, Any], + localns: dict[str, Any], +) -> dict[str, Any]: + """Resolve only annotations declared directly on ``owner``. + + ``typing.get_type_hints(cls)`` merges annotations across the full MRO. + That is wrong for py_class phase 2 because inherited C++/c_class fields + are already registered by the parent type. Resolving only the owner + annotations also avoids evaluating parent annotations in the child + module's namespace. + + Python 3.10's ``typing.get_type_hints`` rejects ``dataclasses.KW_ONLY`` + because it is a singleton, not a type. Replace it in the temporary shim + annotations, then restore it in the resolved hints. + """ + annotations = own_annotations(owner) + if not annotations: + return {} + shim_annotations = dict(annotations) + kw_only_names: list[str] = [] + for name, annotation in annotations.items(): + if _is_kw_only_annotation(annotation): + shim_annotations[name] = _KWOnlyAnnotation + kw_only_names.append(name) + shim = type( + f"_{owner.__name__}OwnAnnotations", + (), + {"__annotations__": shim_annotations, "__module__": owner.__module__}, + ) + kwargs: dict[str, Any] = {"globalns": globalns, "localns": localns} + if sys.version_info >= (3, 11): + kwargs["include_extras"] = True + hints = typing.get_type_hints(shim, **kwargs) + for name in kw_only_names: + hints[name] = KW_ONLY + return hints + + +def resolve_type_hints_by_owner( + cls: type, + globalns: dict[str, Any], +) -> ResolvedFields | None: + """Resolve field annotations grouped by owner class. + + Returns ``(owners, hints_by_owner)`` when every annotation can be resolved. + Returns :data:`None` when a forward reference is still unavailable, signaling + that the caller should defer field registration and retry later. + """ + # Resolve string annotations to types; return None (defer) on NameError. + # + # First try with module-scoped localns (standard Python name resolution). + # On NameError, retry with a cross-module localns that includes classes + # from every registered module — this handles circular imports where the + # target of a forward reference is imported only under TYPE_CHECKING and + # therefore never enters the declaring module's globals. + owners = _field_owner_classes(cls) + localns = _build_localns(cls) + localns.update({owner.__name__: owner for owner in owners}) + try: + hints_by_owner = { + owner: _resolve_own_type_hints( + owner, + getattr(sys.modules.get(owner.__module__, None), "__dict__", globalns), + localns, + ) + for owner in owners + } + except (NameError, AttributeError): + localns = _build_localns(cls, cross_module=True) + localns.update({owner.__name__: owner for owner in owners}) + try: + hints_by_owner = { + owner: _resolve_own_type_hints( + owner, + getattr(sys.modules.get(owner.__module__, None), "__dict__", globalns), + localns, + ) + for owner in owners + } + except (NameError, AttributeError): + return None + return owners, hints_by_owner + + +# --------------------------------------------------------------------------- +# Deferred resolution +# --------------------------------------------------------------------------- + + +def flush_pending() -> None: + """Retry deferred classes until no additional annotations resolve. + + A successful finalization can make another pending class resolvable, so this + function runs to a fixed point. The field-registration finalizer is imported + lazily to keep this module focused on resolution and to avoid an import + cycle with :mod:`tvm_ffi.dataclasses.py_class`. + """ + from .py_class import on_fields_resolved # noqa: PLC0415 + + changed = True + while changed: + changed = False + remaining: list[_PendingClass] = [] + for entry in _PENDING_CLASSES: + resolved = resolve_type_hints_by_owner(entry.cls, entry.globalns) + if resolved is None: + remaining.append(entry) + else: + on_fields_resolved(entry.type_info, resolved) + changed = True + _PENDING_CLASSES[:] = remaining Review Comment:  In the current implementation of `flush_pending`, if `on_fields_resolved` raises an exception for any of the pending classes, the loop is interrupted and `_PENDING_CLASSES` is left unmodified. This means that successfully resolved classes from the current iteration are not removed from `_PENDING_CLASSES`, and the failed class is also not rolled back or removed, leaving the global state inconsistent. We can make this robust by updating `_PENDING_CLASSES` in-place as we resolve each class, and wrapping the finalization in a `try...except` block to roll back the registration of the failed class if an error occurs. ```python def flush_pending() -> None: """Retry deferred classes until no additional annotations resolve. A successful finalization can make another pending class resolvable, so this function runs to a fixed point. The field-registration finalizer is imported lazily to keep this module focused on resolution and to avoid an import cycle with :mod:`tvm_ffi.dataclasses.py_class`. """ from .py_class import on_fields_resolved # noqa: PLC0415 changed = True while changed: changed = False i = 0 while i < len(_PENDING_CLASSES): entry = _PENDING_CLASSES[i] resolved = resolve_type_hints_by_owner(entry.cls, entry.globalns) if resolved is not None: _PENDING_CLASSES.pop(i) try: on_fields_resolved(entry.type_info, resolved) except Exception: rollback_registration(entry.cls, entry.type_info) raise changed = True else: i += 1 ``` ########## include/tvm/ffi/device.h: ########## @@ -0,0 +1,155 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file tvm/ffi/device.h + * \brief Device handling. + */ +#ifndef TVM_FFI_DEVICE_H_ +#define TVM_FFI_DEVICE_H_ + +#include <dlpack/dlpack.h> +#include <tvm/ffi/string.h> +#include <tvm/ffi/type_traits.h> + +#include <cstdint> +#include <limits> +#include <optional> +#include <string> +#include <string_view> + +namespace tvm { +namespace ffi { +namespace details { + +TVM_FFI_INLINE static std::optional<DLDeviceType> TryParseDLDeviceType(std::string_view name) { + if (name == "llvm" || name == "cpu" || name == "c" || name == "test") return kDLCPU; + if (name == "cuda" || name == "nvptx") return kDLCUDA; + if (name == "cl" || name == "opencl") return kDLOpenCL; + if (name == "vulkan") return kDLVulkan; + if (name == "metal") return kDLMetal; + if (name == "vpi") return kDLVPI; + if (name == "rocm") return kDLROCM; + if (name == "ext_dev") return kDLExtDev; + if (name == "hexagon") return kDLHexagon; + if (name == "webgpu") return kDLWebGPU; + if (name == "maia") return kDLMAIA; + if (name == "trn") return kDLTrn; + return std::nullopt; +} + +TVM_FFI_INLINE static std::optional<int32_t> TryParseDLDeviceIndex(std::string_view index) { + if (index.empty()) return std::nullopt; + int64_t sign = 1; + size_t pos = 0; + if (index[0] == '-') { + sign = -1; + pos = 1; + } + if (pos >= index.size()) return std::nullopt; + int64_t value = 0; + for (; pos < index.size(); ++pos) { + char ch = index[pos]; + if (ch < '0' || ch > '9') return std::nullopt; + value = value * 10 + (ch - '0'); + if (value > std::numeric_limits<int32_t>::max()) return std::nullopt; + } + value *= sign; + if (value < std::numeric_limits<int32_t>::min() || value > std::numeric_limits<int32_t>::max()) { + return std::nullopt; + } + return static_cast<int32_t>(value); +} Review Comment:  The current implementation of `TryParseDLDeviceIndex` has a bug when parsing the minimum possible `int32_t` value (`-2147483648`). Because the absolute value `2147483648` is checked against `std::numeric_limits<int32_t>::max()` (`2147483647`) inside the loop before the negative sign is applied, the function incorrectly returns `std::nullopt` for `-2147483648`. We can fix this by allowing the absolute value to reach `max() + 1` during the loop, and then validating the final signed value against the full `int32_t` range after applying the sign. ```c TVM_FFI_INLINE static std::optional<int32_t> TryParseDLDeviceIndex(std::string_view index) { if (index.empty()) return std::nullopt; int64_t sign = 1; size_t pos = 0; if (index[0] == '-') { sign = -1; pos = 1; } if (pos >= index.size()) return std::nullopt; int64_t value = 0; for (; pos < index.size(); ++pos) { char ch = index[pos]; if (ch < '0' || ch > '9') return std::nullopt; value = value * 10 + (ch - '0'); if (value > static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1) return std::nullopt; } value *= sign; if (value < std::numeric_limits<int32_t>::min() || value > std::numeric_limits<int32_t>::max()) { return std::nullopt; } return static_cast<int32_t>(value); } ``` ########## python/tvm_ffi/dataclasses/py_class.py: ########## @@ -521,24 +402,25 @@ def _register_fields_into_type( # ``tvm_ffi.dataclasses.fields()`` compat layer can recover defaults # and default_factory values. _register_fields preserves order, so # own_fields and type_info.fields line up 1:1. + assert len(own_fields) == len(type_info.fields) for py_field, type_field in zip(own_fields, type_info.fields): type_field.dataclass_field = py_field Review Comment:  Since this codebase requires Python 3.10+, we can use `zip(..., strict=True)` instead of a separate `assert` statement. This ensures that a `ValueError` is raised if the lengths of `own_fields` and `type_info.fields` do not match, even when Python is run with optimizations disabled (which removes `assert` statements). ```suggestion for py_field, type_field in zip(own_fields, type_info.fields, strict=True): type_field.dataclass_field = py_field ``` ########## python/tvm_ffi/dataclasses/_resolve_fields.py: ########## @@ -0,0 +1,462 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Forward-annotation resolution for ``@py_class`` registration. + +``@py_class`` registers a Python type before it resolves field annotations, so +self-referential and mutually recursive annotations can refer to classes that +are already present in the FFI type registry. This module owns that early +registration step, the per-module local namespace used by +``typing.get_type_hints``, and the queue of classes whose annotations must be +retried after later definitions become available. + +When a class is ready to finalize, this module lazily imports +``py_class.on_fields_resolved``. Keeping the finalizer import local preserves +the boundary between annotation resolution and field registration while avoiding +a top-level import cycle. +""" + +from __future__ import annotations + +import sys +import typing +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from .. import core +from .field import KW_ONLY + +if TYPE_CHECKING: + from ..core import TypeInfo +else: + TypeInfo = Any + +ResolvedFields = tuple[list[type], dict[type, dict[str, Any]]] + + +# --------------------------------------------------------------------------- +# Module-level state +# --------------------------------------------------------------------------- +# +# ``@py_class`` registration happens in two phases: +# +# Phase 1 (register_type_without_fields) +# Allocates a C-level type index and inserts the class into the +# global type registry. This must happen early so that self- +# referential and mutually-referential annotations can resolve +# the class via ``TypeSchema.from_annotation()``. Phase 1 always +# succeeds or raises immediately for invalid parents. Decorator +# options are stored by py_class.py on ``TypeInfo._decorator_args``. +# +# Phase 2 (py_class.on_fields_resolved) +# Resolves string annotations via ``typing.get_type_hints``, +# returns per-owner resolved hints, and asks py_class.py to +# materialize ``Field`` objects, register fields/methods/type +# attributes with the Cython layer, and install dataclass dunders. +# +# If ``get_type_hints`` raises ``NameError`` (forward reference +# not yet defined), the class is added to ``_PENDING_CLASSES`` +# and retried after each successful phase-2. If phase-2 fails +# for any other reason, ``rollback_registration`` undoes phase-1 +# so the type key can be reused. +# --------------------------------------------------------------------------- + + +@dataclass +class _PendingClass: + """Resolution state for a class waiting on unresolved annotations. + + Only data needed to retry annotation resolution is stored here. Decorator + options live on ``type_info._decorator_args`` so deferred and immediate + finalization use the same source of truth. + """ + + cls: type + type_info: TypeInfo + globalns: dict[str, Any] + + +#: Classes whose phase-2 (field registration) was deferred because +#: ``typing.get_type_hints`` raised ``NameError`` on an unresolved +#: forward reference. Retried after each successful phase-2 via +#: :func:`flush_pending`. +_PENDING_CLASSES: list[_PendingClass] = [] + +#: Per-module mapping of ``class.__name__ → class`` for every +#: ``@py_class``-decorated type. Used as *localns* when resolving +#: annotations so that mutual references between classes in the same +#: module work even before the second class is assigned to the module +#: variable by Python. +_PY_CLASS_BY_MODULE: dict[str, dict[str, type]] = {} + + +class _KWOnlyAnnotation: + """Type-valued stand-in accepted by ``typing.get_type_hints``.""" + + +def _is_kw_only_annotation(annotation: Any) -> bool: + if annotation is KW_ONLY or annotation is _KWOnlyAnnotation: + return True + if isinstance(annotation, str): + return annotation == "KW_ONLY" or annotation.endswith(".KW_ONLY") + forward_arg = getattr(annotation, "__forward_arg__", None) + if isinstance(forward_arg, str): + return forward_arg == "KW_ONLY" or forward_arg.endswith(".KW_ONLY") + return False + + +# --------------------------------------------------------------------------- +# Phase 1: type registration +# --------------------------------------------------------------------------- + + +def _registered_type_info(cls: type) -> TypeInfo | None: + """Return the TypeInfo registered directly for *cls*, not inherited metadata.""" + info = core._type_cls_to_type_info(cls) + if info is not None: + return info + return cls.__dict__.get("__tvm_ffi_type_info__", None) + + +def register_type_without_fields(cls: type, type_key: str | None) -> TypeInfo: + """Register ``cls`` in the FFI type registry before resolving fields. + + The returned :class:`~tvm_ffi.core.TypeInfo` has a type index, type key, and + Python class binding, but no field metadata yet. This early registration is + what lets ``TypeSchema.from_annotation`` resolve references to ``cls`` while + the decorator is still running. + + The class is also inserted into the module-local annotation namespace used + by :func:`resolve_type_hints_by_owner`, allowing sibling ``@py_class`` + declarations in the same module to reference each other before the second + class has been assigned to the module global by Python. + """ + parent_info = next( + (info for base in cls.__mro__[1:] if (info := _registered_type_info(base)) is not None), + None, + ) + if parent_info is None: + raise TypeError( + f"{cls.__name__} must inherit from a registered FFI Object type (e.g. tvm_ffi.Object)" + ) + if type_key is None: + type_key = f"{cls.__module__}.{cls.__qualname__}" + info = core._register_py_class(parent_info, type_key, cls) + setattr(cls, "__tvm_ffi_type_info__", info) + # Register in resolution namespace so sibling classes can find us. + _PY_CLASS_BY_MODULE.setdefault(cls.__module__, {})[cls.__name__] = cls + return info + + +def rollback_registration(cls: type, type_info: TypeInfo) -> None: + """Undo Python-side state from :func:`register_type_without_fields`. + + The C-level type index is permanently consumed (cannot be reclaimed), + but the Python-level registry dicts are cleaned up so a retry with + the same type key does not hit "already registered". The dataclass marker + is also removed because failed phase-2 finalization means the class is not a + usable FFI dataclass. + """ + # Remove from the Cython-level registry dicts (TYPE_KEY_TO_INFO, + # TYPE_CLS_TO_INFO, TYPE_INDEX_TO_INFO, TYPE_INDEX_TO_CLS). + core._rollback_py_class(type_info) # ty: ignore[unresolved-attribute] + # Remove from our own module-level resolution namespace. + _PY_CLASS_BY_MODULE.get(cls.__module__, {}).pop(cls.__name__, None) + for attr in ("__tvm_ffi_type_info__", "__ffi_is_dataclass__"): + try: + delattr(cls, attr) + except AttributeError: + pass + + +# --------------------------------------------------------------------------- +# Annotation resolution +# --------------------------------------------------------------------------- + + +def own_annotations(cls: type) -> dict[str, Any]: + """Return annotations declared directly on ``cls`` without MRO merging.""" + # Python 3.14+ (PEP 749): annotations are lazily evaluated via + # __annotate__ and no longer stored directly in __dict__. getattr() + # triggers evaluation and returns per-class annotations correctly. + # On Python < 3.14, getattr() follows MRO and returns *parent* + # annotations when the child has none — use __dict__ to avoid that. + if sys.version_info >= (3, 14): + return getattr(cls, "__annotations__", {}) + return cls.__dict__.get("__annotations__", {}) + + +def _field_owner_classes(cls: type) -> list[type]: + """Return local MRO entries whose annotations become fields on ``cls``. + + A Python subclass of a registered FFI parent may include unregistered mixin + bases before the registered parent. Those mixin annotations are owned by + the new Python type and should be registered together with ``cls``. Fields + already represented by the nearest registered parent are excluded. + """ + registered_parent = next( + (b for b in cls.__mro__[1:] if _registered_type_info(b) is not None), object + ) + represented = set(registered_parent.__mro__) + return [ + b + for b in reversed(cls.__mro__) + if b is not object and b not in represented and own_annotations(b) + ] + + +def _build_localns(cls: type, *, cross_module: bool = False) -> dict[str, Any]: + """Build the localns dict for resolving ``cls``'s annotations. + + By default, includes only classes from ``cls.__module__``, preserving + standard Python name resolution semantics. When ``cross_module=True``, + also includes classes from all other registered modules as a fallback + — this is needed when ``cls`` has a forward reference to a class in + another module that can't appear in ``cls.__module__``'s globals due + to a circular import (e.g. the target is imported only under + ``if TYPE_CHECKING:``). + + Cross-module entries are added with ``setdefault`` so same-module + classes and the class itself always take precedence over foreign + classes with the same ``__name__``. + """ + localns = dict(_PY_CLASS_BY_MODULE.get(cls.__module__, {})) + localns[cls.__name__] = cls + if cross_module: + for mod_name, mod_classes in list(_PY_CLASS_BY_MODULE.items()): + if mod_name == cls.__module__: + continue + for name, klass in mod_classes.items(): + localns.setdefault(name, klass) + return localns + + +def _resolve_own_type_hints( + owner: type, + globalns: dict[str, Any], + localns: dict[str, Any], +) -> dict[str, Any]: + """Resolve only annotations declared directly on ``owner``. + + ``typing.get_type_hints(cls)`` merges annotations across the full MRO. + That is wrong for py_class phase 2 because inherited C++/c_class fields + are already registered by the parent type. Resolving only the owner + annotations also avoids evaluating parent annotations in the child + module's namespace. + + Python 3.10's ``typing.get_type_hints`` rejects ``dataclasses.KW_ONLY`` + because it is a singleton, not a type. Replace it in the temporary shim + annotations, then restore it in the resolved hints. + """ + annotations = own_annotations(owner) + if not annotations: + return {} + shim_annotations = dict(annotations) + kw_only_names: list[str] = [] + for name, annotation in annotations.items(): + if _is_kw_only_annotation(annotation): + shim_annotations[name] = _KWOnlyAnnotation + kw_only_names.append(name) + shim = type( + f"_{owner.__name__}OwnAnnotations", + (), + {"__annotations__": shim_annotations, "__module__": owner.__module__}, + ) + kwargs: dict[str, Any] = {"globalns": globalns, "localns": localns} + if sys.version_info >= (3, 11): + kwargs["include_extras"] = True + hints = typing.get_type_hints(shim, **kwargs) + for name in kw_only_names: + hints[name] = KW_ONLY + return hints + + +def resolve_type_hints_by_owner( + cls: type, + globalns: dict[str, Any], +) -> ResolvedFields | None: + """Resolve field annotations grouped by owner class. + + Returns ``(owners, hints_by_owner)`` when every annotation can be resolved. + Returns :data:`None` when a forward reference is still unavailable, signaling + that the caller should defer field registration and retry later. + """ + # Resolve string annotations to types; return None (defer) on NameError. + # + # First try with module-scoped localns (standard Python name resolution). + # On NameError, retry with a cross-module localns that includes classes + # from every registered module — this handles circular imports where the + # target of a forward reference is imported only under TYPE_CHECKING and + # therefore never enters the declaring module's globals. + owners = _field_owner_classes(cls) + localns = _build_localns(cls) + localns.update({owner.__name__: owner for owner in owners}) + try: + hints_by_owner = { + owner: _resolve_own_type_hints( + owner, + getattr(sys.modules.get(owner.__module__, None), "__dict__", globalns), + localns, + ) + for owner in owners + } + except (NameError, AttributeError): + localns = _build_localns(cls, cross_module=True) + localns.update({owner.__name__: owner for owner in owners}) + try: + hints_by_owner = { + owner: _resolve_own_type_hints( + owner, + getattr(sys.modules.get(owner.__module__, None), "__dict__", globalns), + localns, + ) + for owner in owners + } + except (NameError, AttributeError): + return None + return owners, hints_by_owner + + +# --------------------------------------------------------------------------- +# Deferred resolution +# --------------------------------------------------------------------------- + + +def flush_pending() -> None: + """Retry deferred classes until no additional annotations resolve. + + A successful finalization can make another pending class resolvable, so this + function runs to a fixed point. The field-registration finalizer is imported + lazily to keep this module focused on resolution and to avoid an import + cycle with :mod:`tvm_ffi.dataclasses.py_class`. + """ + from .py_class import on_fields_resolved # noqa: PLC0415 + + changed = True + while changed: + changed = False + remaining: list[_PendingClass] = [] + for entry in _PENDING_CLASSES: + resolved = resolve_type_hints_by_owner(entry.cls, entry.globalns) + if resolved is None: + remaining.append(entry) + else: + on_fields_resolved(entry.type_info, resolved) + changed = True + _PENDING_CLASSES[:] = remaining + + +def _raise_unresolved_forward_reference(cls: type, globalns: dict[str, Any]) -> None: + """Raise :class:`TypeError` listing the annotations that cannot be resolved.""" + localns = _build_localns(cls, cross_module=True) + owners = _field_owner_classes(cls) + localns.update({owner.__name__: owner for owner in owners}) + unresolved: list[str] = [] + for owner in owners: + for name, ann_str in own_annotations(owner).items(): + if isinstance(ann_str, str): + try: + eval(ann_str, globalns, localns) # pylint: disable=eval-used + except (NameError, AttributeError) as err: + unresolved.append(f"{name}: {ann_str} ({err})") Review Comment:  When diagnosing unresolved forward references, `eval` is used to check which annotations are failing. However, evaluating arbitrary annotation strings can raise exceptions other than `NameError` or `AttributeError` (such as `SyntaxError`, `TypeError`, or `ValueError`). Catching `Exception` ensures that any evaluation failure is gracefully captured and reported in the final `TypeError` message rather than crashing the diagnostic routine itself. ```suggestion try: eval(ann_str, globalns, localns) # pylint: disable=eval-used except Exception as err: unresolved.append(f"{name}: {ann_str} ({err})") ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
