Skip to content

Validate function argument annotations by name#1633

Open
fallintoplace wants to merge 2 commits into
NVIDIA:mainfrom
fallintoplace:agent/validate-function-annotations
Open

Validate function argument annotations by name#1633
fallintoplace wants to merge 2 commits into
NVIDIA:mainfrom
fallintoplace:agent/validate-function-annotations

Conversation

@fallintoplace

@fallintoplace fallintoplace commented Jul 9, 2026

Copy link
Copy Markdown

Summary

  • Validate required parameter annotations by name across code generation, replay helpers, kernel overloads, and JAX FFI callables.
  • Parse arguments in declaration order while preserving return annotations for function return validation.
  • Add regressions for kernels, functions, overload stubs, and JAX callables.

Why

The old count-based validation treated a return annotation as an argument annotation. A function such as f(a: int, b) -> None therefore passed validation and failed later with a less useful codegen or FFI error.

Validation

  • Ruff lint and format checks pass for all changed Python files.
  • Targeted codegen regression passes with the CPU-only native runtime.
  • Targeted JAX FFI regression passes with CPU JAX 0.5+.
  • Python compilation and git diff --check pass.

Summary by CodeRabbit

  • Bug Fixes
    • Strengthened validation for incomplete function/kernel/JAX FFI argument type annotations.
    • Improved errors to name the first specific missing parameter.
    • Preserved positional argument order when processing kernels, overloads, and JAX FFI callables.
    • Tightened return annotation handling to require None where applicable.
  • Tests
    • Added new test coverage for incomplete argument annotations across codegen, overloads, and JAX interoperability.

Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented Jul 9, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Jul 9, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Enterprise

Run ID: 9ac1f89d-44a0-4113-8dfe-2c9c388a7c56

📥 Commits

Reviewing files that changed from the base of the PR and between 0479f7c and eae79b2.

📒 Files selected for processing (2)
  • warp/_src/context.py
  • warp/tests/test_codegen.py
🚧 Files skipped from review as they are similar to previous changes (2)
  • warp/tests/test_codegen.py
  • warp/_src/context.py

📝 Walkthrough

Walkthrough

The change strengthens missing-annotation and return-annotation validation across Warp codegen, overloads, and JAX FFI callables. Argument processing now follows declared parameter order, with tests covering codegen and JAX error reporting.

Changes

Annotation validation

Layer / File(s) Summary
Core annotation validation and ordering
warp/_src/context.py, warp/_src/codegen.py, warp/tests/test_codegen.py
Missing parameters are identified by name, overload return annotations are validated, adjoint variables follow declaration order, and codegen tests cover the resulting errors.
JAX FFI annotation validation
warp/_src/jax/ffi.py, warp/tests/interop/test_jax.py
JAX FFI callables validate argument and return annotations, preserve argument order, and test incomplete annotations.

Estimated code review effort: 3 (Moderate) | ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 18.75% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately summarizes the main change: validating function argument annotations by parameter name.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@greptile-apps

greptile-apps Bot commented Jul 9, 2026

Copy link
Copy Markdown

Greptile Summary

This PR replaces the previous count-based annotation validation (len(annotations) < len(args)) with a name-based check that correctly identifies which specific parameter is missing a type annotation. It also refactors iteration in codegen.py, context.py, and jax/ffi.py to use argspec.args (declaration order) instead of iterating over argspec.annotations (which included the "return" key and required explicit skipping).

  • context.py get_function_args: now returns {name: annotations[name] for name in argspec.args}, properly excluding the "return" key that was previously leaked to callers (notably the replay-function path at line 1448–1449).
  • codegen.py arg-building loop: replaced for name, type in adj.arg_types.items() (with a if name == "return": continue guard) with for name in argspec.args:, producing the same result more cleanly and guaranteeing declaration order.
  • jax/ffi.py: the return-type validity check is now a single early guard before the iteration loop, and iteration over argspec.args avoids the need to skip the "return" key mid-loop. Tests for incomplete annotations are added across all four code paths.

Confidence Score: 5/5

Safe to merge — the change is a targeted validation improvement with no mutations to codegen logic or FFI dispatch.

All three code paths (codegen, overload, JAX FFI) apply the same name-based check consistently. The iteration refactor from adj.arg_types.items() to argspec.args removes the previously needed return-key guard without changing any output. get_function_args now correctly excludes the return key, which directly fixes a latent replay-helper overload-lookup failure. Test coverage spans every touched code path.

No files require special attention.

Important Files Changed

Filename Overview
warp/_src/codegen.py Replaces count-based annotation check with name-based list comprehension; arg-building loop now iterates argspec.args instead of adj.arg_types.items(), cleanly removing the need for the return key guard.
warp/_src/context.py get_function_args now returns only parameter annotations (no return key), fixing the replay-helper overload-lookup bug; overload stub validation similarly switched to name-based check with argspec.args iteration.
warp/_src/jax/ffi.py Missing-annotation check and return-type guard are now both early exits before the arg-building loop; iteration over argspec.args eliminates the return key skip and preserves declaration order.
warp/tests/test_codegen.py Adds test_error_incomplete_argument_annotations covering get_function_args, @wp.kernel, @wp.func, and @wp.overload paths; also verifies that a fully-annotated-with-return function correctly yields only arg annotations.
warp/tests/interop/test_jax.py Adds test_ffi_jax_callable_incomplete_argument_annotations covering the JAX FFI path; registered with devices=None since it validates at construction time rather than execution.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Function / Kernel decorated] --> B[get_full_arg_spec]
    B --> C{argspec.args vs argspec.annotations}
    C -->|name-based check| D["missing = [name for name in argspec.args if name not in argspec.annotations]"]
    D --> E{missing list empty?}
    E -->|No| F["Raise: Argument 'X' in function 'Y' must be type annotated"]
    E -->|Yes| G[Return-type guard: if return_type is not None]
    G -->|Non-None return| H[Raise TypeError]
    G -->|None or absent| I[Iterate argspec.args in order]
    I --> J[Build arg list / FfiArg list using argspec.annotations per name]
    J --> K[Proceed to codegen / JAX FFI setup]

    style F fill:#f66,color:#fff
    style H fill:#f66,color:#fff
    style K fill:#6a6,color:#fff
Loading
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A[Function / Kernel decorated] --> B[get_full_arg_spec]
    B --> C{argspec.args vs argspec.annotations}
    C -->|name-based check| D["missing = [name for name in argspec.args if name not in argspec.annotations]"]
    D --> E{missing list empty?}
    E -->|No| F["Raise: Argument 'X' in function 'Y' must be type annotated"]
    E -->|Yes| G[Return-type guard: if return_type is not None]
    G -->|Non-None return| H[Raise TypeError]
    G -->|None or absent| I[Iterate argspec.args in order]
    I --> J[Build arg list / FfiArg list using argspec.annotations per name]
    J --> K[Proceed to codegen / JAX FFI setup]

    style F fill:#f66,color:#fff
    style H fill:#f66,color:#fff
    style K fill:#6a6,color:#fff
Loading

Reviews (2): Last reviewed commit: "Exclude return annotations from replay a..." | Re-trigger Greptile

Comment thread warp/_src/context.py Outdated

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
warp/tests/interop/test_jax.py (1)

864-872: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Consider adding a positive-case test for -> None-annotated, fully-typed functions.

This test only covers the missing-annotation path. Given the concern raised in warp/_src/jax/ffi.py (lines 569-572) about explicit -> None return annotations, a companion test asserting that a fully-annotated function with -> None succeeds would guard against regressing that previously-fixed behavior.

Do you want me to draft this additional test case?

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@warp/tests/interop/test_jax.py` around lines 864 - 872, The JAX FFI tests
lack coverage confirming that fully annotated functions with an explicit None
return annotation are accepted. Add a positive companion test near
test_ffi_jax_callable_incomplete_argument_annotations defining a function whose
arguments use valid Warp annotations and whose return type is -> None, then call
wp.jax_callable with num_outputs=1 and assert it succeeds without raising.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@warp/tests/interop/test_jax.py`:
- Around line 864-872: The JAX FFI tests lack coverage confirming that fully
annotated functions with an explicit None return annotation are accepted. Add a
positive companion test near
test_ffi_jax_callable_incomplete_argument_annotations defining a function whose
arguments use valid Warp annotations and whose return type is -> None, then call
wp.jax_callable with num_outputs=1 and assert it succeeds without raising.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Enterprise

Run ID: 14e9face-7caf-483c-8b59-e2ba996e4405

📥 Commits

Reviewing files that changed from the base of the PR and between 8667276 and 0479f7c.

📒 Files selected for processing (5)
  • warp/_src/codegen.py
  • warp/_src/context.py
  • warp/_src/jax/ffi.py
  • warp/tests/interop/test_jax.py
  • warp/tests/test_codegen.py

Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant