Return types for array subscript annotations (GH-1548)#1549
Return types for array subscript annotations (GH-1548)#1549FabienPean-Virtonomy wants to merge 4 commits into
Conversation
`wp.array[...]`-style subscripts previously produced lightweight *instances*, which can’t participate in PEP 604 unions (e.g. `wp.array[float] | float`) and forces per-use allocations. Return a cached lightweight annotation *type* instead, while preserving the codegen-facing metadata (`dtype`, `ndim`, `vars`, `__ctype__`) and updating the subscript-type tests to match. Signed-off-by: Fabien Péan <pean@virtonomy.io>
|
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 (4)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThis PR refactors Warp's array annotation system to return parameterized type objects (via a metaclass-based factory) instead of lightweight instances, updates detection and signature encoding to operate on those types, adjusts a couple runtime call sites, updates tests, and documents the change in the changelog. ChangesArray Annotation Type System
🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
warp/_src/types.py (1)
5132-5153:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReturning classes here breaks the existing
None-argument inference path.After this change,
wp.array[...]is a class object, butinfer_argument_types()still reconstructs array templates withtype(t)(dtype=t.dtype, ndim=t.ndim)on Line 7163. For these annotations,type(t)is now_ArrayAnnotationTypeMeta, so passingNonefor an array-typed parameter will raise instead of preserving the template type.At minimum, the downstream reconstruction needs to special-case
is_array_annotation(t)and reuset(or rebuild it via_make_array_annotation_type(...)) instead of callingtype(t)(...).🤖 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/_src/types.py` around lines 5132 - 5153, The change made in _parse_array_subscript returns class-like objects which breaks infer_argument_types' reconstruction that calls type(t)(dtype=..., ndim=...) (type(t) is now _ArrayAnnotationTypeMeta); update infer_argument_types to special-case array annotations by using is_array_annotation(t) and either reuse the existing annotation object t directly when filling in None arguments or rebuild the annotation via _make_array_annotation_type(_ARRAY_ANNOTATION_MAP.get(base), dtype=t.dtype, ndim=t.ndim) instead of calling type(t)(...), ensuring array templates preserve their original semantics.
🤖 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.
Inline comments:
In `@warp/_src/types.py`:
- Around line 5112-5129: _make_array_annotation_type currently constructs a new
_ArrayAnnotationTypeMeta on every call causing identity instability; add a cache
(e.g., a module-level dict or an attribute on ann_base) keyed by the normalized
dtype and ndim (use the same dtype = dtype if dtype is Any else
type_to_warp(dtype) logic and the resolved ndim) and return the cached class
when present instead of allocating a new one; update the function to compute the
cache key (ann_base, dtype, ndim), check and return from cache, and only create
and store a new _ArrayAnnotationTypeMeta if the key is missing so repeated
subscripts (e.g. wp.array[float]) yield identical types.
- Around line 5064-5070: The __ctype__ implementation for array-like classes
currently returns a bare array_t() which discards annotated rank/shape/strides;
change __ctype__ (using cls._concrete_cls, array_t(), indexedarray_t,
ARRAY_MAX_DIMS) to preserve the annotated ndim/shape/strides the same way the
old annotation path did—read cls.ndim (falling back to Any), cls.shape and
cls.strides if present (or default placeholders) and construct array_t(...) with
those values instead of returning array_t() so codegen/native consumers retain
the array rank information.
---
Outside diff comments:
In `@warp/_src/types.py`:
- Around line 5132-5153: The change made in _parse_array_subscript returns
class-like objects which breaks infer_argument_types' reconstruction that calls
type(t)(dtype=..., ndim=...) (type(t) is now _ArrayAnnotationTypeMeta); update
infer_argument_types to special-case array annotations by using
is_array_annotation(t) and either reuse the existing annotation object t
directly when filling in None arguments or rebuild the annotation via
_make_array_annotation_type(_ARRAY_ANNOTATION_MAP.get(base), dtype=t.dtype,
ndim=t.ndim) instead of calling type(t)(...), ensuring array templates preserve
their original semantics.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Enterprise
Run ID: 25e54390-be3e-49bd-94f0-acb5c563f40b
📒 Files selected for processing (3)
CHANGELOG.mdwarp/_src/types.pywarp/tests/test_subscript_types.py
Greptile SummaryThis PR changes
Confidence Score: 5/5Safe to merge — the change is well-scoped, the cache correctly ensures identity equality, and all call sites that previously relied on instance checks have been updated to the new type-based helpers. The implementation is mechanically sound: the metaclass correctly intercepts all class-level operations, _make_array_annotation_type caches by (ann_base, dtype, ndim) so repeated subscriptions are free, and is_array_annotation correctly gates on the presence of dtype/ndim/_concrete_cls to exclude unparameterized base types. The builtins.py fix (concrete_array_type instead of type()) is necessary and correct. The only findings are a stale docstring phrase and a redundant line in ctype. No files require special attention. Important Files Changed
Reviews (2): Last reviewed commit: "Fix array annotation type handling to co..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def __eq__(cls, other): | ||
| if not isinstance(other, type): | ||
| return NotImplemented | ||
| if not issubclass(other, _ArrayAnnotationBase): | ||
| return False |
There was a problem hiding this comment.
__eq__ implementation returns False (not NotImplemented) when other is a type but is not an _ArrayAnnotationBase subclass. Python convention is to return NotImplemented when a comparison is not meaningful for the operand types, so the other side can still be tried. Returning False suppresses that fallback. In practice this rarely matters here (built-in type.__eq__ also falls back correctly), but it diverges from the standard contract and could silently short-circuit user-defined __eq__ on custom type objects compared against annotation types.
| def __eq__(cls, other): | |
| if not isinstance(other, type): | |
| return NotImplemented | |
| if not issubclass(other, _ArrayAnnotationBase): | |
| return False | |
| def __eq__(cls, other): | |
| if not isinstance(other, type): | |
| return NotImplemented | |
| if not issubclass(other, _ArrayAnnotationBase): | |
| return NotImplemented |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
/ok to test 2b76761 |
Signed-off-by: Fabien Péan <pean@virtonomy.io>
Signed-off-by: Fabien Péan <pean@virtonomy.io>
… method Signed-off-by: Fabien Péan <pean@virtonomy.io>
|
Thanks for the fix. It works and the tests are thorough. The root cause is just that |
2 participants
Description
wp.array[...]-style subscripts previously produced lightweight instances, which cannot participate in PEP 604 unions (e.g.wp.array[float] | float) and forces per-use allocations. Return a cached lightweight annotation type instead, while preserving the codegen-facing metadata (dtype,ndim,vars,__ctype__) and updating the subscript-type tests to match.Closes #1548
Checklist
Unreleasedsection.Validation summary
Tests updated and added to
test_subscript_types.pytest suiteBug fix
Summary by CodeRabbit
Bug Fixes
Documentation