mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 15:37:44 +00:00
fix(auto-fixer): fix chained groupby(level=N).groupby('date') pattern
LLM learns from feedback to use groupby(level=1) for instrument, then
chains .groupby('date') to add the date dimension — but DataFrameGroupBy
has no .groupby() method, causing AttributeError at runtime.
Replace the invalid chain with a correct two-level groupby using
index.get_level_values(), consistent with the existing instrument+date fix.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -56,7 +56,8 @@ class FactorAutoFixer:
|
||||
self._fix_reset_index_groupby, # First: fix groupby(level=N) after reset_index()
|
||||
self._fix_groupby_mixed_levels, # Second: fix groupby(level=[int, str])
|
||||
self._fix_groupby_column_on_multiindex, # Third: fix groupby(['instrument','date']) on MultiIndex
|
||||
self._fix_rolling_ddof, # Fourth: remove unsupported ddof kwarg
|
||||
self._fix_chained_groupby, # Fourth: fix groupby(level=N).groupby('date') chain
|
||||
self._fix_rolling_ddof, # Fifth: remove unsupported ddof kwarg
|
||||
self._fix_groupby_apply_to_transform, # Fifth: fix groupby patterns
|
||||
self._fix_min_periods, # Sixth: fix min_periods in rolling calls
|
||||
self._fix_inf_nan_handling, # Seventh: add inf/nan handling
|
||||
@@ -177,6 +178,37 @@ class FactorAutoFixer:
|
||||
|
||||
return fixed_code
|
||||
|
||||
def _fix_chained_groupby(self, code: str) -> str:
|
||||
"""
|
||||
Fix: var.groupby(level=N).groupby('date') is invalid — DataFrameGroupBy has no
|
||||
.groupby() method. The LLM learns this pattern from poisoned feedback that told
|
||||
it to use groupby(level=1) for instrument, then tries to add a date dimension on
|
||||
top.
|
||||
|
||||
Replace the entire chain with a proper two-level groupby:
|
||||
var.groupby(level=1).groupby('date')
|
||||
→ var.groupby([var.index.get_level_values(1),
|
||||
var.index.get_level_values(0).normalize()])
|
||||
"""
|
||||
fixed_code = code
|
||||
|
||||
def _replace_chained(m: re.Match) -> str:
|
||||
var = m.group(1)
|
||||
repl = (
|
||||
f"{var}.groupby([{var}.index.get_level_values(1), "
|
||||
f"{var}.index.get_level_values(0).normalize()])"
|
||||
)
|
||||
self.fixes_applied.append(f"chained_groupby: {m.group(0)[:60]} → two-level")
|
||||
return repl
|
||||
|
||||
# var.groupby(level=N).groupby('date') or .groupby("date")
|
||||
fixed_code = re.sub(
|
||||
r'(\w+)\.groupby\(level=\d+\)\.groupby\(["\']date["\']\)',
|
||||
_replace_chained,
|
||||
fixed_code,
|
||||
)
|
||||
return fixed_code
|
||||
|
||||
def _fix_rolling_ddof(self, code: str) -> str:
|
||||
"""
|
||||
Fix: pandas rolling() does not accept a ddof kwarg — raises TypeError.
|
||||
|
||||
@@ -58,6 +58,21 @@ class TestGroupbyColumnOnMultiindex:
|
||||
assert "groupby('instrument')" in result
|
||||
|
||||
|
||||
class TestChainedGroupby:
|
||||
def test_chained_groupby_level_then_date(self, fixer):
|
||||
code = "df.groupby(level=1).groupby('date')['price_volume'].transform('cumsum')"
|
||||
result = fixer.fix(code)
|
||||
assert "get_level_values(1)" in result
|
||||
assert "get_level_values(0).normalize()" in result
|
||||
assert ".groupby('date')" not in result
|
||||
|
||||
def test_chained_groupby_with_double_quotes(self, fixer):
|
||||
code = 'df.groupby(level=0).groupby("date")["col"].sum()'
|
||||
result = fixer.fix(code)
|
||||
assert "get_level_values" in result
|
||||
assert '.groupby("date")' not in result
|
||||
|
||||
|
||||
class TestRollingDdof:
|
||||
def test_removes_ddof_from_rolling_args(self, fixer):
|
||||
result = fixer.fix("df.rolling(20, min_periods=1, ddof=1).std()")
|
||||
|
||||
Reference in New Issue
Block a user