feat(python): add @py_class decorator for Python-defined FFI dataclasses

[junrushao]

Summary

Add @py_class, a new decorator that lets users define TVM FFI Object types entirely in Python with dataclass-style field annotations. This complements @c_class (which wraps C++-backed types) by enabling pure-Python FFI types that participate in the same object system -- type registry, reference counting, cross-language serialization, and packed-function interop.

Motivation

Currently, defining a new FFI-visible type requires either writing C++ (ObjectDef<T>) or using @c_class which expects C++ reflection metadata. @py_class removes the C++ requirement entirely: annotate fields with Python types, and the decorator handles type-index allocation, field registration, and dunder generation.

Usage

from tvm_ffi.dataclasses import py_class, field
from tvm_ffi import Object

@py_class
class Point(Object):
    x: float
    y: float

@py_class("my.namespace.Line", eq=True)
class Line(Object):
    start: Point
    end: Point
    label: str = "default"

Architecture

Registration is split into two phases to handle forward and mutual references:

• Phase 1 (_phase1_register_type): Allocates a C-level type index and inserts the class into the global type registry. Runs eagerly so that self-referential and mutually-referential annotations can be resolved.
• 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.

When get_type_hints raises NameError (forward reference not yet defined), the class is deferred to a pending list and retried after each successful phase-2. A temporary __init__ triggers lazy completion on first instantiation. Failed phase-2 rolls back phase-1 so the type key can be reused.

Public Interfaces

Symbol	Location	Description
@py_class	tvm_ffi.dataclasses.py_class	Main decorator (bare, string-key, and kwarg calling conventions)
Parameters	type_key, init, repr, eq, order, unsafe_hash, kw_only, slots	Mirrors dataclasses.dataclass semantics

Test Plan

• 246 tests across 47 test classes (tests/python/test_dataclass_py_class.py, 3533 lines)
• Coverage includes: basic registration, field parsing, defaults/factories, kw_only, ClassVar, init generation, post_init, repr, equality, ordering, hashing, copy/deepcopy, single/multi-level inheritance, forward/mutual references, dunder preservation, registration rollback, type conversion errors, getter/setter corner cases, memory lifetime, bool alignment, optional fields, Any fields, FFI global function interop
• All 1817 Python tests pass (25 skipped, 3 xfailed) -- no regressions
• All pre-commit hooks pass

Untested Edge Cases

• Thread safety of the pending-class flush mechanism
• Interaction with importlib.reload() on modules containing @py_class-decorated types
• Deeply nested generic type annotations (e.g. List[Optional[Map[str, Array[int]]]])

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the Python FFI system by introducing a new decorator that allows developers to define FFI classes directly in Python using a familiar dataclass-like syntax. This change simplifies the process of creating and managing Python-backed FFI types, improving developer experience and promoting more idiomatic Python code within the TVM ecosystem. The new decorator handles complex aspects like type resolution and dunder method generation automatically, reducing boilerplate and potential errors.

Highlights

• New @py_class decorator: Introduced the @py_class decorator, enabling the definition of TVM FFI classes entirely in Python with dataclass-like semantics. This includes typed fields, automatic __init__, structural __eq__/__hash__/__repr__, and two-phase registration for handling forward and self-referential types.
• Exported FFI dataclass components: Exported py_class, Field, field, and KW_ONLY from tvm_ffi.dataclasses to provide a comprehensive API for defining Python-backed FFI types.
• Comprehensive test suite: Added a comprehensive test suite of 246 tests specifically for the py_class decorator, covering its functionality, field API, inheritance, type conversion, copy/hash behavior, and various edge cases.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces the @py_class decorator, enabling the definition of TVM FFI dataclasses directly in Python. This is a significant enhancement, providing dataclass-like semantics including typed fields, __init__ generation, structural comparison, hashing, and robust handling of forward/self-referential types through a two-phase registration mechanism. The implementation demonstrates careful consideration for edge cases, such as KW_ONLY fields, user-defined dunder preservation, and error handling with registration rollback. The accompanying test suite is exceptionally comprehensive, covering a wide array of scenarios from basic field parsing and inheritance to complex type conversions and memory management, ensuring the reliability and correctness of this new feature.

[junrushao]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the @py_class decorator, a significant feature that enables defining TVM FFI objects purely in Python with a dataclass-like syntax. The implementation is robust, handling complex scenarios like forward and mutual references through a two-phase registration process and lazy initialization. The code is well-structured and documented. I've identified a critical thread-safety issue in the registration logic that should be addressed, as well as a minor opportunity for code simplification. Overall, this is an excellent contribution.