mitsuba-renderer / drjit

Dr.Jit — A Just-In-Time-Compiler for Differentiable Rendering
BSD 3-Clause "New" or "Revised" License
572 stars 40 forks source link

Handle postponed type information in dataclasses #249

Closed njroussel closed 1 month ago

njroussel commented 1 month ago

This PR adds handling for postponed type hints (PEP 563) on Dr.Jit functions involving a dataclass. These are particularly useful when using Mitsuba, as a dataclass will most likely be defined before the variant has been set. Here's an example:

from __future__ import annotations
from dataclasses import dataclass

import mitsuba as mi
import drjit as dr

@dataclass
class SampleData:
    data: mi.Float # postponed type annotation (invalid before mi.set_variant())

def main():
    sample_data = dr.alloc_local(SampleData, 10) # at this point in time `mi.Float` can be resolved

if __name__ == "__main__":
    mi.set_variant("cuda_ad_rgb")
    main()

The implementation relies on typing.get_type_hints, which is the recommended way of resolving type hints.