Validate function argument annotations by name#1633
Conversation
Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThe 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. ChangesAnnotation validation
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Greptile SummaryThis PR replaces the previous count-based annotation validation (
Confidence Score: 5/5Safe 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
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
%%{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
Reviews (2): Last reviewed commit: "Exclude return annotations from replay a..." | Re-trigger Greptile |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
warp/tests/interop/test_jax.py (1)
864-872: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider 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-> Nonereturn annotations, a companion test asserting that a fully-annotated function with-> Nonesucceeds 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
📒 Files selected for processing (5)
warp/_src/codegen.pywarp/_src/context.pywarp/_src/jax/ffi.pywarp/tests/interop/test_jax.pywarp/tests/test_codegen.py
Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>
Summary
Why
The old count-based validation treated a return annotation as an argument annotation. A function such as
f(a: int, b) -> Nonetherefore passed validation and failed later with a less useful codegen or FFI error.Validation
git diff --checkpass.Summary by CodeRabbit
Nonewhere applicable.