[REQ] Support for `ast.AnnAssign` (type-annotated assignments) in kernels

[liblaf]

Parent Issue: https://github.com/NVIDIA/warp/issues/549

Description

Currently, using type-annotated assignments within a @wp.func or @wp.kernel decorated function raises a WarpCodegenError.

For instance, the following code will fail:

@wp.func
def add(a: wp.vec3, b: wp.vec3) -> wp.vec3:
    c: wp.vec3 = a + b
    return c

The error produced is: WarpCodegenError: Construct 'ast.AnnAssign' not supported in kernels.

This feature request is to add support for ast.AnnAssign to allow developers to use type annotations for local variables within Warp kernels.

Context

Type annotations for variables in Python are designed for static analysis and do not have a runtime impact. They are ignored by the Python runtime and are primarily intended to assist developers and third-party tools like IDEs and type checkers. Quoting from typing — Support for type hints — Python 3.13.5 documentation:

The Python runtime does not enforce function and variable type annotations. They can be used by third party tools such as type checkers, IDEs, linters, etc.

By supporting ast.AnnAssign, Warp would allow for more readable and self-documenting code. These annotations are a significant benefit for developers, code reviewers, and static analysis tools in understanding the code's intent without affecting performance.

Since Python's runtime does not enforce type annotations, the Warp compiler could treat an annotated assignment such as c: wp.vec3 = a + b as a standard assignment (c = a + b) during the code generation process.