132 lines
4.3 KiB
Python
132 lines
4.3 KiB
Python
from __future__ import annotations
|
|
import tokenize
|
|
import io
|
|
import os
|
|
import sys
|
|
|
|
def format_code(source: str) -> str:
|
|
"""
|
|
Formats Python code to use exactly 1 space for indentation (including continuations),
|
|
max 1 blank line between top-level definitions, and 0 blank lines inside
|
|
function/method bodies.
|
|
"""
|
|
if not source:
|
|
return ""
|
|
try:
|
|
tokens = list(tokenize.generate_tokens(io.StringIO(source).readline))
|
|
except tokenize.TokenError:
|
|
return source # Return as-is if it's not valid python (e.g. template files)
|
|
lines = source.splitlines(keepends=True)
|
|
num_lines = len(lines)
|
|
block_level = 0
|
|
paren_level = 0
|
|
in_function_stack = []
|
|
expecting_function_indent = False
|
|
line_indent = {}
|
|
line_is_blank = {i: True for i in range(1, num_lines + 2)}
|
|
line_is_string_interior = {i: False for i in range(1, num_lines + 2)}
|
|
line_seen = set()
|
|
pending_blank_lines = []
|
|
for tok in tokens:
|
|
t_type = tok.type
|
|
t_string = tok.string
|
|
start_line, _ = tok.start
|
|
end_line, _ = tok.end
|
|
if t_type == tokenize.STRING:
|
|
for l in range(start_line + 1, end_line + 1):
|
|
line_is_string_interior[l] = True
|
|
if t_type not in (tokenize.NL, tokenize.NEWLINE, tokenize.INDENT, tokenize.DEDENT, tokenize.ENDMARKER):
|
|
for l in range(start_line, end_line + 1):
|
|
line_is_blank[l] = False
|
|
pending_blank_lines = []
|
|
if t_type == tokenize.INDENT:
|
|
block_level += 1
|
|
if expecting_function_indent:
|
|
in_function_stack.append(block_level)
|
|
expecting_function_indent = False
|
|
elif t_type == tokenize.DEDENT:
|
|
block_level -= 1
|
|
if in_function_stack and block_level < in_function_stack[-1]:
|
|
in_function_stack.pop()
|
|
for l in pending_blank_lines:
|
|
line_indent[l] = block_level + paren_level
|
|
if t_string in (')', ']', '}'):
|
|
paren_level -= 1
|
|
if start_line not in line_seen:
|
|
line_indent[start_line] = block_level + paren_level
|
|
if t_type not in (tokenize.INDENT, tokenize.DEDENT):
|
|
line_seen.add(start_line)
|
|
if t_type in (tokenize.NL, tokenize.NEWLINE):
|
|
pending_blank_lines.append(start_line)
|
|
if t_type == tokenize.NAME and t_string == 'def':
|
|
expecting_function_indent = True
|
|
if t_string in ('(', '[', '{'):
|
|
paren_level += 1
|
|
output = []
|
|
consecutive_blanks = 0
|
|
for i in range(1, num_lines + 1):
|
|
if line_is_string_interior[i]:
|
|
output.append(lines[i-1])
|
|
continue
|
|
if line_is_blank[i]:
|
|
indent = line_indent.get(i, 0)
|
|
if indent > 0:
|
|
continue
|
|
else:
|
|
if consecutive_blanks < 1:
|
|
output.append("\n")
|
|
consecutive_blanks += 1
|
|
continue
|
|
original_line = lines[i-1]
|
|
indent = line_indent.get(i, 0)
|
|
stripped = original_line.lstrip()
|
|
is_def_start = stripped.startswith(('def ', 'class ', 'async def ', '@'))
|
|
if is_def_start and output and consecutive_blanks == 0:
|
|
prev_line = output[-1].strip()
|
|
if prev_line and not prev_line.endswith(':') and not prev_line.startswith('@'):
|
|
output.append("\n")
|
|
consecutive_blanks += 1
|
|
consecutive_blanks = 0
|
|
output.append(" " * indent + stripped)
|
|
if not stripped.endswith('\n') and i < num_lines:
|
|
output[-1] += '\n'
|
|
if output and not output[-1].endswith('\n'):
|
|
output[-1] += '\n'
|
|
return "".join(output)
|
|
|
|
def process_file(file_path: str, write: bool) -> None:
|
|
try:
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
formatted = format_code(content)
|
|
if write:
|
|
if formatted != content:
|
|
with open(file_path, "w", encoding="utf-8") as f:
|
|
f.write(formatted)
|
|
print(f"Formatted: {file_path}")
|
|
else:
|
|
sys.stdout.reconfigure(encoding='utf-8')
|
|
sys.stdout.write(formatted)
|
|
except Exception as e:
|
|
print(f"Error processing {file_path}: {e}")
|
|
|
|
def main() -> None:
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description="AI-optimized Python code formatter.")
|
|
parser.add_argument("paths", nargs="+", help="Files or directories to format.")
|
|
parser.add_argument("--write", action="store_true", help="Write changes back to files.")
|
|
parser.add_argument("--exclude", nargs="*", default=[".venv", "__pycache__", ".git"], help="Directories to exclude.")
|
|
args = parser.parse_args()
|
|
for path in args.paths:
|
|
if os.path.isfile(path):
|
|
process_file(path, args.write)
|
|
elif os.path.isdir(path):
|
|
for root, dirs, files in os.walk(path):
|
|
dirs[:] = [d for d in dirs if d not in args.exclude]
|
|
for file in files:
|
|
if file.endswith(".py"):
|
|
process_file(os.path.join(root, file), args.write)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|