1866: fix: use DataFrame.__name__ instead of hardcoded string in metaclass method wrapping - #73
1866: fix: use DataFrame.__name__ instead of hardcoded string in metaclass method wrapping#73martin-augment wants to merge 4 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (6)
WalkthroughTwo independent changes are bundled in this PR. In ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request replaces hardcoded "DataFrame" string checks with DataFrame.name in ballista/extension.py and adds unit tests to verify that methods returning DataFrame are correctly wrapped. The review feedback suggests a more defensive approach when accessing annotations directly on callables to prevent potential AttributeError exceptions, particularly for C-extension methods.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if ( | ||
| callable(base_value) | ||
| and not base_name.startswith("__") | ||
| and base_value.__annotations__.get("return") == "DataFrame" | ||
| and base_value.__annotations__.get("return") == DataFrame.__name__ | ||
| ): |
There was a problem hiding this comment.
Accessing __annotations__ directly on a callable can raise an AttributeError if the callable (such as a built-in or C-extension method from PyO3) does not define it or if it is None. Using getattr(base_value, "__annotations__", None) or {} is a safer, more defensive approach to prevent potential runtime crashes.
| if ( | |
| callable(base_value) | |
| and not base_name.startswith("__") | |
| and base_value.__annotations__.get("return") == "DataFrame" | |
| and base_value.__annotations__.get("return") == DataFrame.__name__ | |
| ): | |
| if ( | |
| callable(base_value) | |
| and not base_name.startswith("__") | |
| and (getattr(base_value, "__annotations__", None) or {}).get("return") == DataFrame.__name__ | |
| ): |
There was a problem hiding this comment.
value:annoying; category:bug; feedback: The Gemini AI reviewer is not correct! The DataFrame class uses type hints and all its members are static, i.e. there are no dynamically added ones. Raising an AttributeError and failing the test would be better than silently omit the assertion for a new field that is not typed.
| if ( | ||
| callable(base_value) | ||
| and not base_name.startswith("__") | ||
| and base_value.__annotations__.get("return") == "DataFrame" | ||
| and base_value.__annotations__.get("return") == DataFrame.__name__ | ||
| ): |
There was a problem hiding this comment.
Accessing __annotations__ directly on a callable can raise an AttributeError if the callable (such as a built-in or C-extension method from PyO3) does not define it or if it is None. Using getattr(base_value, "__annotations__", None) or {} is a safer, more defensive approach to prevent potential runtime crashes.
| if ( | |
| callable(base_value) | |
| and not base_name.startswith("__") | |
| and base_value.__annotations__.get("return") == "DataFrame" | |
| and base_value.__annotations__.get("return") == DataFrame.__name__ | |
| ): | |
| if ( | |
| callable(base_value) | |
| and not base_name.startswith("__") | |
| and (getattr(base_value, "__annotations__", None) or {}).get("return") == DataFrame.__name__ | |
| ): |
There was a problem hiding this comment.
value:annoying; category:bug; feedback: The Gemini AI reviewer is not correct! The DataFrame class uses type hints and all its members are static, i.e. there are no dynamically added ones. Raising an AttributeError and failing the test would be better than silently omit the assertion for a new field that is not typed.
| should_be_wrapped = { | ||
| name | ||
| for name, val in base_cls.__dict__.items() | ||
| if callable(val) | ||
| and not name.startswith("__") | ||
| and val.__annotations__.get("return") == DataFrame.__name__ | ||
| } |
There was a problem hiding this comment.
Accessing __annotations__ directly on a callable can raise an AttributeError if the callable does not define it or if it is None. Using getattr(val, "__annotations__", None) or {} is a safer, more defensive approach to prevent potential runtime crashes.
| should_be_wrapped = { | |
| name | |
| for name, val in base_cls.__dict__.items() | |
| if callable(val) | |
| and not name.startswith("__") | |
| and val.__annotations__.get("return") == DataFrame.__name__ | |
| } | |
| should_be_wrapped = { | |
| name | |
| for name, val in base_cls.__dict__.items() | |
| if callable(val) | |
| and not name.startswith("__") | |
| and (getattr(val, "__annotations__", None) or {}).get("return") == DataFrame.__name__ | |
| } |
There was a problem hiding this comment.
value:annoying; category:bug; feedback: The Gemini AI reviewer is not correct! The DataFrame class uses type hints and all its members are static, i.e. there are no dynamically added ones. Raising an AttributeError and failing the test would be better than silently omit the assertion for a new field that is not typed.
Code Review:
|
| Area | Verdict |
|---|---|
| Core annotation change | Acceptable — marginal maintainability improvement |
| Leftover blank comment lines | Minor — clean them up |
| New tests | Good — thorough structural coverage |
| Cargo.lock dependency additions | Needs explanation |
AI config files (CLAUDE.md, AGENTS.md, .cursor/rules.md, .gemini/rules.md) |
Remove before merging — prompt injection |
Recommendation: Do not merge as-is. Strip the AI configuration files from the branch. The Python change itself (annotation comparison + tests) is reasonable once the blank comment lines are tidied.
🤖 Augment PR SummarySummary: This PR refactors how Ballista’s Python metaclasses detect “DataFrame-returning” methods so the check uses the Changes:
Technical Notes: The wrapping logic is driven purely by method return annotations, so upstream annotation format changes can affect which methods get wrapped. 🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| callable(base_value) | ||
| and not base_name.startswith("__") | ||
| and base_value.__annotations__.get("return") == "DataFrame" | ||
| and base_value.__annotations__.get("return") == DataFrame.__name__ |
There was a problem hiding this comment.
python/python/ballista/extension.py:79: This __annotations__["return"] == DataFrame.__name__ check only matches string return annotations equal to the class name; if DataFusion ever emits an actual DataFrame type (or a qualified string), these methods would silently stop being wrapped and return plain DataFrame instances.
Severity: medium
Other Locations
python/python/ballista/extension.py:118python/python/tests/test_context.py:150
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:annoying; category:bug; feedback: The Augment AI reviewer is not correct! If the return type is not "DataFrame" then the unit tests will fail and the respective changes will be made to fix the problem. Until then there is no problem to be fixed.
| for name, val in base_cls.__dict__.items() | ||
| if callable(val) | ||
| and not name.startswith("__") | ||
| and val.__annotations__.get("return") == DataFrame.__name__ |
There was a problem hiding this comment.
python/python/tests/test_context.py:150: _assert_dataframe_returning_methods_wrapped assumes every callable(val) has a __annotations__ dict; if any DataFusion method is a C-extension/builtin descriptor without __annotations__, this will raise AttributeError and fail the test for reasons unrelated to wrapping.
Severity: low
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:annoying; category:bug; feedback: The Augment AI reviewer is not correct! The DataFrame class uses type hints and all its members are static, i.e. there are no dynamically added ones. Raising an AttributeError and failing the test would be better than silently omit the assertion for a new field that is not typed.
value:useful; category:bug; feedback: The Claude AI reviewer is correct! There are no changes in Cargo.toml, so no changes are expected in Cargo.lock too. None of the added dependencies are related to the rest of the changes in the Pull Request |
3 participants
1866: To review by AI