fix(auto-fixer): add groupby([level=N,'date']) SyntaxError fix

LLM generates invalid Python by putting keyword args inside lists:
  df.groupby([level=1, 'date'])  ← SyntaxError

Also fixes the regex for the chained groupby Pattern A/B which had
an unescaped ')' causing re.error that silently reverted the fix.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
TPTBusiness
2026-04-26 21:01:22 +02:00
parent ba4d64b434
commit ff1c9fc554
2 changed files with 46 additions and 12 deletions
@@ -182,33 +182,53 @@ class FactorAutoFixer:
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.
Fix two broken patterns the LLM generates when trying to group by (instrument, date):
Replace the entire chain with a proper two-level groupby:
Pattern A — chained groupby (runtime AttributeError):
var.groupby(level=1).groupby('date')
→ var.groupby([var.index.get_level_values(1),
var.index.get_level_values(0).normalize()])
Pattern B — keyword arg inside list (SyntaxError):
var.groupby([level=1, 'date'])
→ same two-level replacement
"""
fixed_code = code
def _replace_chained(m: re.Match) -> str:
var = m.group(1)
repl = (
def _two_level(var: str, tag: str) -> str:
self.fixes_applied.append(f"chained_groupby: {tag} → two-level")
return (
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")
# Pattern A: var.groupby(level=N).groupby('date')
fixed_code = re.sub(
r'(\w+)\.groupby\(level=\d+\)\.groupby\(["\']date["\']\)',
_replace_chained,
lambda m: _two_level(m.group(1), m.group(0)[:60]),
fixed_code,
)
# Pattern B: .groupby([level=N, 'date']) — SyntaxError in Python.
# The variable before .groupby may be complex (e.g. df[mask]) so we don't
# try to capture it; we use df as the index reference (always correct since
# all filtered frames share df's MultiIndex structure).
def _two_level_df(tag: str) -> str:
self.fixes_applied.append(f"chained_groupby: {tag} → two-level")
return ".groupby([df.index.get_level_values(1), df.index.get_level_values(0).normalize()])"
fixed_code = re.sub(
r'\.groupby\(\[\s*level\s*=\s*\d+\s*,\s*["\']?date["\']?\s*\]\)',
lambda m: _two_level_df(m.group(0)[:60]),
fixed_code,
)
# Also handle reversed order: ['date', level=N]
fixed_code = re.sub(
r'\.groupby\(\[\s*["\']?date["\']?\s*,\s*level\s*=\s*\d+\s*\]\)',
lambda m: _two_level_df(m.group(0)[:60]),
fixed_code,
)
return fixed_code
def _fix_rolling_ddof(self, code: str) -> str:
+14
View File
@@ -72,6 +72,20 @@ class TestChainedGroupby:
assert "get_level_values" in result
assert '.groupby("date")' not in result
def test_list_with_level_keyword_syntax_error(self, fixer):
# groupby([level=1, 'date']) is a SyntaxError — must be fixed before execution
code = "asian_vol = df[mask].groupby([level=1, 'date'])['log_return'].std()"
result = fixer.fix(code)
assert "get_level_values(1)" in result
assert "normalize()" in result
assert "level=1," not in result
def test_list_with_level_keyword_reversed(self, fixer):
code = "df.groupby(['date', level=1])['x'].mean()"
result = fixer.fix(code)
assert "get_level_values" in result
assert "level=1" not in result
class TestMinPeriodsNotTouched:
def test_small_min_periods_preserved(self, fixer):