gemini-code-assist[bot] commented on code in PR #506: URL: https://github.com/apache/tvm-ffi/pull/506#discussion_r2967577588
########## 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:  The `try...except AttributeError` blocks around `delattr` are redundant. The attribute's existence is checked before attempting deletion (`if name in cls.__dict__` and `if class_val is not MISSING`), so `delattr` should not raise an `AttributeError`. Removing these blocks simplifies the code and improves readability. ```suggestion if name in cls.__dict__: delattr(cls, name) 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: delattr(cls, name) ``` ########## 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:  The module-level state `_PENDING_CLASSES` and `_PY_CLASS_BY_MODULE` are accessed and modified from multiple places without synchronization, which can lead to race conditions in a multi-threaded environment (e.g., if classes are defined dynamically in different threads). This was noted as an untested edge case in the PR description. To ensure thread safety, a `threading.Lock` should be used to protect all accesses to this shared state. I suggest the following changes: 1. Add a lock at the module level: ```python import sys import threading # ... other imports _REGISTRATION_LOCK = threading.Lock() ``` 2. Acquire the lock within the main `decorator` function in `py_class`: ```python def decorator(cls: _T) -> _T: with _REGISTRATION_LOCK: nonlocal effective_type_key # ... rest of the function ``` 3. Use double-checked locking in the temporary `__init__` to handle concurrent instantiations safely and efficiently: ```python # in _make_temporary_init() def __init__(self: Any, *args: Any, **kwargs: Any) -> None: if type_info.fields is None: with _REGISTRATION_LOCK: if type_info.fields is None: # Check again inside lock try: if not _phase2_register_fields(cls, type_info, globalns, params): _raise_unresolved_forward_reference(cls, globalns) _flush_pending() except Exception: _PENDING_CLASSES[:] = [p for p in _PENDING_CLASSES if p.cls is not cls] _rollback_registration(cls, type_info) raise cls.__init__(self, *args, **kwargs) ``` -- 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]
