junrushao commented on code in PR #506: URL: https://github.com/apache/tvm-ffi/pull/506#discussion_r2967598565
########## python/tvm_ffi/dataclasses/py_class.py: ########## @@ -0,0 +1,459 @@ +# 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. +"""The ``py_class`` decorator: Python-defined FFI classes with dataclass semantics.""" + +from __future__ import annotations + +import sys +import typing +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, ClassVar, TypeVar + +from typing_extensions import dataclass_transform + +from .. import core +from ..core import MISSING, TypeSchema +from ..registry import _add_class_attrs, _install_dataclass_dunders +from .field import KW_ONLY, Field, field + +_T = TypeVar("_T", bound=type) + + +# --------------------------------------------------------------------------- +# Module-level state +# --------------------------------------------------------------------------- +# +# Registration happens in two phases: +# +# Phase 1 (_phase1_register_type) +# 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 non-Object parents). +# +# Phase 2 (_phase2_register_fields) +# Resolves string annotations via ``typing.get_type_hints``, +# converts them to ``TypeSchema`` / ``Field`` objects, validates +# field ordering, registers fields with the Cython layer, and +# installs ``__init__``, ``__repr__``, ``__eq__``, etc. +# +# 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: + """Bookkeeping for a class whose annotations couldn't be resolved yet.""" + + cls: type + type_info: Any # core.TypeInfo + globalns: dict[str, Any] + params: 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]] = {} Review Comment: Thanks for the review! This is a valid observation. However, Python's GIL already serializes bytecode execution, and `@py_class` decoration happens at module-import time which is inherently single-threaded per module. Adding `threading.Lock` + double-checked locking would add complexity for a scenario that doesn't arise in practice (dynamically defining `@py_class` classes from multiple threads simultaneously). This is already called out explicitly in the PR description under "Untested Edge Cases." We can revisit if a real use case emerges. ########## python/tvm_ffi/dataclasses/py_class.py: ########## @@ -0,0 +1,459 @@ +# 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. +"""The ``py_class`` decorator: Python-defined FFI classes with dataclass semantics.""" + +from __future__ import annotations + +import sys +import typing +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, ClassVar, TypeVar + +from typing_extensions import dataclass_transform + +from .. import core +from ..core import MISSING, TypeSchema +from ..registry import _add_class_attrs, _install_dataclass_dunders +from .field import KW_ONLY, Field, field + +_T = TypeVar("_T", bound=type) + + +# --------------------------------------------------------------------------- +# Module-level state +# --------------------------------------------------------------------------- +# +# Registration happens in two phases: +# +# Phase 1 (_phase1_register_type) +# 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 non-Object parents). +# +# Phase 2 (_phase2_register_fields) +# Resolves string annotations via ``typing.get_type_hints``, +# converts them to ``TypeSchema`` / ``Field`` objects, validates +# field ordering, registers fields with the Cython layer, and +# installs ``__init__``, ``__repr__``, ``__eq__``, etc. +# +# 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: + """Bookkeeping for a class whose annotations couldn't be resolved yet.""" + + cls: type + type_info: Any # core.TypeInfo + globalns: dict[str, Any] + params: 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]] = {} + + +# --------------------------------------------------------------------------- +# Phase 1: type registration +# --------------------------------------------------------------------------- + + +def _phase1_register_type(cls: type, type_key: str | None) -> Any: + """Phase 1: allocate type index and register the type (always succeeds).""" + parent_info: core.TypeInfo | None = None + for base in cls.__bases__: + parent_info = core._type_cls_to_type_info(base) + if parent_info is not None: + break + 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: Any) -> None: + """Undo :func:`_phase1_register_type` after a phase-2 failure. + + 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". + """ + # 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) + try: + delattr(cls, "__tvm_ffi_type_info__") + except AttributeError: + pass + + +# --------------------------------------------------------------------------- +# Phase 2: annotation resolution, field registration, dunder installation +# --------------------------------------------------------------------------- + + +def _collect_own_fields( + cls: type, + hints: dict[str, Any], + decorator_kw_only: bool, +) -> list[Field]: + """Parse own annotations into :class:`Field` objects. + + - Skips ``ClassVar`` annotations. + - Handles ``KW_ONLY`` sentinel. + - Extracts ``Field`` metadata from class attributes (set via :func:`field`). + - Handles bare defaults (non-``Field`` values). + - Converts resolved types to ``TypeSchema``. + - Resolves ``hash=None`` to follow ``compare``. + """ + fields: list[Field] = [] + kw_only_active = decorator_kw_only + own_annotations: dict[str, str] = getattr(cls, "__annotations__", {}) + + for name in own_annotations: + resolved_type = hints.get(name) + # Skip ClassVar + if ( + resolved_type is None + or resolved_type is ClassVar + or typing.get_origin(resolved_type) is ClassVar + ): + continue + + # KW_ONLY sentinel + if resolved_type is KW_ONLY: + kw_only_active = True + if name in cls.__dict__: + try: + delattr(cls, name) + except AttributeError: + pass + continue + + # Extract Field from class dict (inline of _pop_field_from_class) + class_val = cls.__dict__.get(name, MISSING) + if isinstance(class_val, Field): + f = class_val + elif class_val is not MISSING: + f = field(default=class_val) + else: + f = field() + if class_val is not MISSING: + try: + delattr(cls, name) + except AttributeError: + pass Review Comment: Good catch on the redundancy. The `try/except AttributeError` is a defensive pattern against exotic metaclasses or descriptors with custom `__delete__` that might raise unexpectedly. It's only 3 extra lines per site and harmless, so I'd prefer to keep it for robustness. Removing it saves little but introduces a small risk under edge-case metaclass usage. -- 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]
