import re
def fix_file(filepath):
with open(filepath, 'r') as f:
lines = f.readlines()
new_lines = []
i = 0
fixed = 0
total = len(lines)
while i < total:
line = lines[i]
stripped = line.strip()
if stripped == 'conn = None':
new_lines.append(line)
i += 1
continue
if not re.match(r'^\s*conn\s*=\s*get_db\s*\(\s*\)', line):
new_lines.append(line)
i += 1
continue
indent = len(line) - len(line.lstrip())
indent_str = ' ' * indent
# Check already wrapped (conn=None + try: before get_db)
found_wrap = False
for wi in range(max(0, i - 6), i):
if lines[wi].strip() == 'conn = None':
for tj in range(wi + 1, i):
if lines[tj].strip() == 'try:':
found_wrap = True
break
if found_wrap:
break
if found_wrap:
new_lines.append(line)
i += 1
continue
# ── Pattern B: try: immediately on next line ──
if i + 1 < total and lines[i + 1].strip() == 'try:':
try_indent = len(lines[i + 1]) - len(lines[i + 1].lstrip())
finally_idx = None
scan_end = min(i + 200, total)
for sj in range(i + 2, scan_end):
s = lines[sj].strip()
if s == '':
continue
sj_indent = len(lines[sj]) - len(lines[sj].lstrip())
if sj_indent < try_indent:
break
if sj_indent == try_indent:
if s == 'finally:':
finally_idx = sj
break
elif s.startswith('except'):
# After except, keep scanning for finally:
# (try/except/finally is valid)
continue
else:
# Other line at try_indent without finally — block ended
break
if finally_idx is not None:
# Find end of finally body
finally_body_end = finally_idx + 1
while (finally_body_end < total and
lines[finally_body_end].strip() != '' and
len(lines[finally_body_end]) - len(lines[finally_body_end].lstrip()) > try_indent):
finally_body_end += 1
new_lines.append(f'{indent_str}conn = None\n')
new_lines.append(f'{indent_str}try:\n')
new_lines.append(f'{indent_str} conn = get_db()\n')
for bk in range(i + 2, finally_idx):
new_lines.append(lines[bk])
new_lines.append(f'{indent_str}finally:\n')
new_lines.append(f'{indent_str} if conn:\n')
new_lines.append(f'{indent_str} conn.close()\n')
i = finally_body_end
fixed += 1
continue
else:
# No finally — check for conn.close() after try/except
close_idx = None
for ci in range(i + 2, min(i + 200, total)):
s = lines[ci].strip()
if s and (s.startswith('def ') or s.startswith('class ')):
break
if s == 'conn.close()':
ci_indent = len(lines[ci]) - len(lines[ci].lstrip())
if ci_indent == indent:
close_idx = ci
break
if close_idx is not None:
new_lines.append(f'{indent_str}conn = None\n')
new_lines.append(f'{indent_str}try:\n')
new_lines.append(f'{indent_str} conn = get_db()\n')
for bk in range(i + 2, close_idx):
new_lines.append(lines[bk])
new_lines.append(f'{indent_str}finally:\n')
new_lines.append(f'{indent_str} if conn:\n')
new_lines.append(f'{indent_str} conn.close()\n')
i = close_idx + 1
fixed += 1
continue
else:
new_lines.append(line)
i += 1
continue
# ── Pattern A: find conn.close() at same indent ──
close_idx = None
for ci in range(i + 1, total):
if lines[ci].strip() == 'conn.close()':
ci_indent = len(lines[ci]) - len(lines[ci].lstrip())
if ci_indent == indent:
close_idx = ci
break
if close_idx is None:
new_lines.append(line)
i += 1
continue
# Check if there's an existing try/finally between get_db() and close()
existing_try_idx = None
existing_finally_idx = None
for ti in range(i + 1, close_idx):
if lines[ti].strip() == 'try:' and len(lines[ti]) - len(lines[ti].lstrip()) == indent:
depth = 1
fi = ti + 1
while fi < close_idx and depth > 0:
s = lines[fi].strip()
if s == 'try:':
depth += 1
elif s.startswith('except') or s.startswith('finally'):
depth -= 1
fi += 1
for fj in range(ti + 1, fi):
if lines[fj].strip() == 'finally:' and len(lines[fj]) - len(lines[fj].lstrip()) == indent:
for fc in range(fj + 1, close_idx + 1):
if lines[fc].strip() == 'conn.close()':
existing_try_idx = ti
existing_finally_idx = fj
break
if existing_try_idx:
break
if existing_try_idx:
break
if existing_try_idx is not None:
# Existing try/finally — add conn=None, move get_db() inside try
finally_body_end = existing_finally_idx + 1
while (finally_body_end < total and
lines[finally_body_end].strip() != '' and
len(lines[finally_body_end]) - len(lines[finally_body_end].lstrip()) > indent):
finally_body_end += 1
new_lines.append(f'{indent_str}conn = None\n')
new_lines.append(lines[existing_try_idx])
for bk in range(i, existing_try_idx):
if lines[bk].strip():
new_lines.append(' ' + lines[bk])
else:
new_lines.append(lines[bk])
for bk in range(existing_try_idx + 1, existing_finally_idx):
new_lines.append(lines[bk])
new_lines.append(lines[existing_finally_idx])
new_lines.append(f'{indent_str} if conn:\n')
new_lines.append(f'{indent_str} conn.close()\n')
i = finally_body_end
fixed += 1
else:
# No existing try/finally — simple wrap
new_lines.append(f'{indent_str}conn = None\n')
new_lines.append(f'{indent_str}try:\n')
new_lines.append(f'{indent_str} {stripped}\n')
for bk in range(i + 1, close_idx):
new_lines.append(' ' + lines[bk])
new_lines.append(f'{indent_str}finally:\n')
new_lines.append(f'{indent_str} if conn:\n')
new_lines.append(f'{indent_str} conn.close()\n')
i = close_idx + 1
fixed += 1
with open(filepath, 'w') as f:
f.writelines(new_lines)
return fixed
files = [
'app/models.py',
'app/routes/admin.py',
'app/routes/agent_api.py',
'app/routes/auth.py',
'app/routes/email.py',
'app/routes/invites.py',
'app/routes/user_sites.py',
]
total_fixed = 0
for fp in files:
count = fix_file(fp)
if count:
print(f'{fp}: {count} fixes')
else:
print(f'{fp}: 0 fixes')
total_fixed += count
print(f'\nTotal: {total_fixed} fixes applied')