code formater from llm

main.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "libcst",
# ]
# ///

import argparse
import bisect
import builtins
import concurrent.futures
import difflib
import fnmatch
import sys
import tomllib
import tkinter as tk
from tkinter import messagebox
from enum import Enum, auto
from pathlib import Path
from typing import List, Optional, Sequence, Set, Dict

import libcst as cst
import libcst.matchers as m

# ==============================================================================
# CONSTANTS & DEFAULTS
# ==============================================================================

BUILTINS_AND_KEYWORDS = set(dir(builtins)) | {
    'True',
    'False',
    'None',
    'self',
    'cls',
    '__name__',
    '__doc__',
    '__file__',
    '__package__',
}

DEFAULT_EXCLUDE = [
    '.git',
    '__pycache__',
    'venv',
    'env',
    '.venv',
    '.tox',
    'build',
    'dist',
    '.mypy_cache',
    '.pytest_cache',
    '*/migrations/*',
]


# ==============================================================================
# STATUS ENUMS
# ==============================================================================


class FileStatus(Enum):
    UNCHANGED = auto()
    CHANGED = auto()
    ERROR = auto()


class BlockType(Enum):
    DOCSTRING = auto()
    IMPORT = auto()
    CONSTANT = auto()
    OTHER = auto()
    CLASS = auto()
    FUNCTION = auto()
    MAIN = auto()


# ==============================================================================
# CONFIGURATION
# ==============================================================================


class Config:
    def __init__(self):
        self.alphabetical: bool = False
        self.sort_imports: bool = False
        self.check: bool = False
        self.diff: bool = False
        self.relaxed: bool = False
        self.verbose: bool = False
        self.quiet: bool = False
        self.exclude: List[str] = DEFAULT_EXCLUDE.copy()


def load_pyproject_toml() -> dict:
    path = Path('pyproject.toml')
    if path.is_file():
        try:
            with open(path, 'rb') as f:
                data = tomllib.load(f)
                return data.get('tool', {}).get('funkshway', {})
        except Exception as e:
            print(f'Warning: Failed to parse pyproject.toml: {e}', file=sys.stderr)
    return {}


# ==============================================================================
# CORE LOGIC
# ==============================================================================


class UnrecognizedBlockError(Exception):
    pass


_DUMMY_MODULE = cst.Module(body=[])


class Block:
    def __init__(self, type: BlockType, nodes: List[cst.CSTNode], index: int):
        self.type = type
        self.nodes = nodes
        self.index = index
        self.provides: Set[str] = set()
        self.requires: Set[str] = set()
        self._lexical_key: Optional[str] = None

    @property
    def lexical_key(self) -> str:
        if self._lexical_key is None:
            self._lexical_key = ''.join(_DUMMY_MODULE.code_for_node(n) for n in self.nodes).strip()
        return self._lexical_key

    def get_sort_name(self) -> str:
        if self.provides:
            return min(self.provides).lower()
        return self.lexical_key.lower()

    def __repr__(self) -> str:
        return f'Block(type={self.type}, provides={self.provides}, requires={self.requires})'


def topo_sort(blocks: List[Block], reverse: bool = False, key_func=None) -> List[Block]:
    if not blocks:
        return []

    if key_func is None:
        key_func = lambda b: b.lexical_key

    group_set = set(blocks)
    all_provides: Dict[str, Block] = {}
    for b in blocks:
        for p in b.provides:
            all_provides[p] = b

    adj: Dict[Block, List[Block]] = {b: [] for b in blocks}
    in_degree: Dict[Block, int] = {b: 0 for b in blocks}

    for b in blocks:
        for r in b.requires:
            if r in all_provides:
                provider = all_provides[r]
                if provider != b:
                    u, v = (provider, b) if not reverse else (b, provider)
                    if u in group_set and v in group_set:
                        adj[u].append(v)
                        in_degree[v] += 1

    unvisited = set(blocks)
    ready = [b for b in blocks if in_degree[b] == 0]
    ready.sort(key=key_func)

    result = []
    while len(result) < len(blocks):
        if ready:
            next_b = ready.pop(0)
        else:
            cycle_candidates = sorted(list(unvisited), key=key_func)
            next_b = cycle_candidates[0]

        result.append(next_b)
        unvisited.remove(next_b)

        for neighbor in adj[next_b]:
            if neighbor in unvisited:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    bisect.insort(ready, neighbor, key=key_func)

    return result


def sort_blocks(blocks: List[Block], alphabetical: bool = False, sort_imports: bool = False) -> List[Block]:
    by_type = {t: [] for t in BlockType}
    for b in blocks:
        by_type[b.type].append(b)

    sorted_blocks = []
    for t in [
        BlockType.DOCSTRING,
        BlockType.IMPORT,
        BlockType.CONSTANT,
        BlockType.OTHER,
        BlockType.CLASS,
        BlockType.FUNCTION,
        BlockType.MAIN,
    ]:
        group = by_type[t]
        if not group:
            continue

        if t == BlockType.IMPORT:
            if sort_imports:
                for b in group:
                    b.nodes.sort(key=lambda n: _DUMMY_MODULE.code_for_node(n).strip())
            sorted_blocks.extend(group)
            continue

        if alphabetical and t in (BlockType.CLASS, BlockType.CONSTANT, BlockType.FUNCTION):
            group.sort(key=lambda b: b.get_sort_name())
        elif t in (BlockType.CONSTANT, BlockType.CLASS):
            group = topo_sort(group, reverse=False, key_func=lambda b: b.lexical_key)
        elif t == BlockType.OTHER:
            group = topo_sort(group, reverse=False, key_func=lambda b: b.index)
        elif t == BlockType.FUNCTION:
            dunders = [
                b for b in group if b.provides and any(p.startswith('__') and p.endswith('__') for p in b.provides)
            ]
            others = [b for b in group if b not in dunders]

            dunders.sort(key=lambda b: b.lexical_key)
            others = topo_sort(others, reverse=True, key_func=lambda b: b.lexical_key)
            group = dunders + others
        else:
            group.sort(key=lambda b: b.lexical_key)

        sorted_blocks.extend(group)
    return sorted_blocks


def is_constant_assignment(node: cst.CSTNode) -> bool:
    if m.matches(
        node,
        m.SimpleStatementLine(
            body=[m.Assign(targets=[m.AssignTarget(target=m.Name(value=m.MatchRegex(r'^[A-Z_][A-Z0-9_]*$')))])]
        ),
    ):
        return True
    if m.matches(
        node, m.SimpleStatementLine(body=[m.AnnAssign(target=m.Name(value=m.MatchRegex(r'^[A-Z_][A-Z0-9_]*$')))])
    ):
        return True
    if m.matches(node, m.TypeAlias()):
        return True
    return False


def is_main_block(node: cst.CSTNode) -> bool:
    if m.matches(
        node,
        m.If(
            test=m.Comparison(
                left=m.Name(value='__name__'),
                comparisons=[m.ComparisonTarget(operator=m.Equal(), comparator=m.SimpleString(value='"__main__"'))],
            )
        ),
    ):
        return True
    if m.matches(
        node,
        m.If(
            test=m.Comparison(
                left=m.Name(value='__name__'),
                comparisons=[m.ComparisonTarget(operator=m.Equal(), comparator=m.SimpleString(value="'__main__'"))],
            )
        ),
    ):
        return True
    return False


def get_block_type(node: cst.CSTNode) -> BlockType:
    if m.matches(node, m.SimpleStatementLine(body=[m.Import() | m.ImportFrom()])):
        return BlockType.IMPORT
    if is_constant_assignment(node):
        return BlockType.CONSTANT
    if m.matches(node, m.ClassDef()):
        return BlockType.CLASS
    if m.matches(node, m.FunctionDef()):
        return BlockType.FUNCTION
    if is_main_block(node):
        return BlockType.MAIN
    if m.matches(node, m.SimpleStatementLine(body=[m.Expr(value=m.SimpleString())])):
        return BlockType.DOCSTRING
    return BlockType.OTHER


class LocalNameCollector(cst.CSTVisitor):
    def __init__(self):
        self.locals: Set[str] = set()

    def visit_Param(self, node: cst.Param) -> None:
        self.locals.add(node.name.value)

    def visit_AssignTarget(self, node: cst.AssignTarget) -> None:
        if isinstance(node.target, cst.Name):
            self.locals.add(node.target.value)

    def visit_AnnAssign(self, node: cst.AnnAssign) -> None:
        if isinstance(node.target, cst.Name):
            self.locals.add(node.target.value)

    def visit_NamedExpr(self, node: cst.NamedExpr) -> None:
        if isinstance(node.target, cst.Name):
            self.locals.add(node.target.value)

    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
        self.locals.add(node.name.value)

    def visit_ClassDef(self, node: cst.ClassDef) -> None:
        self.locals.add(node.name.value)

    def visit_WithItem(self, node: cst.WithItem) -> None:
        # ИСПРАВЛЕНИЕ: В WithItem переменная хранится в .asname, а не .target
        if node.asname:
            for name_node in m.findall(node.asname.name, m.Name()):
                self.locals.add(name_node.value)

    def visit_For(self, node: cst.For) -> None:
        for name_node in m.findall(node.target, m.Name()):
            self.locals.add(name_node.value)

    def visit_CompFor(self, node: cst.CompFor) -> None:
        for name_node in m.findall(node.target, m.Name()):
            self.locals.add(name_node.value)

    def visit_ExceptHandler(self, node: cst.ExceptHandler) -> None:
        if node.name and isinstance(node.name, cst.AsName):
            if isinstance(node.name.name, cst.Name):
                self.locals.add(node.name.name.value)


class TopLevelCollector(cst.CSTVisitor):
    def __init__(self):
        self.provides: Set[str] = set()
        self.requires: Set[str] = set()

    def add_require(self, name_val: str) -> None:
        if name_val not in BUILTINS_AND_KEYWORDS:
            self.requires.add(name_val)

    def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:
        self.provides.add(node.name.value)
        for base in node.bases:
            if isinstance(base.value, cst.Name):
                self.add_require(base.value.value)
        return False

    def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:
        self.provides.add(node.name.value)
        collector = LocalNameCollector()
        node.visit(collector)
        for name_node in m.findall(node.body, m.Name()):
            if name_node.value not in collector.locals:
                self.add_require(name_node.value)
        return False

    def visit_Assign(self, node: cst.Assign) -> Optional[bool]:
        for target in node.targets:
            for name_node in m.findall(target.target, m.Name()):
                self.provides.add(name_node.value)
        node.value.visit(self)
        return False

    def visit_AnnAssign(self, node: cst.AnnAssign) -> Optional[bool]:
        for name_node in m.findall(node.target, m.Name()):
            self.provides.add(name_node.value)
        if node.value:
            node.value.visit(self)
        return False

    def visit_TypeAlias(self, node: cst.TypeAlias) -> Optional[bool]:
        if isinstance(node.name, cst.Name):
            self.provides.add(node.name.value)
        node.value.visit(self)
        return False

    def visit_Attribute(self, node: cst.Attribute) -> Optional[bool]:
        node.value.visit(self)
        return False

    def visit_Arg(self, node: cst.Arg) -> Optional[bool]:
        node.value.visit(self)
        return False

    def visit_ImportAlias(self, node: cst.ImportAlias) -> Optional[bool]:
        if node.asname:
            if isinstance(node.asname.name, cst.Name):
                self.provides.add(node.asname.name.value)
        else:
            if isinstance(node.name, cst.Name):
                self.provides.add(node.name.value)
        return False

    def visit_AsName(self, node: cst.AsName) -> Optional[bool]:
        return False

    def visit_Annotation(self, node: cst.Annotation) -> Optional[bool]:
        return False

    def visit_Name(self, node: cst.Name) -> None:
        self.add_require(node.value)


def group_into_blocks(body: Sequence[cst.BaseStatement]) -> List[Block]:
    blocks: List[Block] = []
    for index, statement in enumerate(body):
        btype = get_block_type(statement)

        if btype == BlockType.IMPORT and blocks and blocks[-1].type == BlockType.IMPORT:
            blocks[-1].nodes.append(statement)
        else:
            block = Block(btype, [statement], index)
            collector = TopLevelCollector()
            statement.visit(collector)
            block.provides.update(collector.provides)
            block.requires.update(collector.requires - collector.provides)
            blocks.append(block)
    return blocks


def format_code(source_code: str, alphabetical: bool = False, sort_imports: bool = False) -> str:
    module = cst.parse_module(source_code)
    blocks = group_into_blocks(module.body)

    if not blocks:
        return source_code

    sorted_blocks = sort_blocks(blocks, alphabetical, sort_imports)

    for i, block in enumerate(sorted_blocks):
        if not block.nodes:
            continue

        first_node = block.nodes[0]
        if not hasattr(first_node, 'leading_lines'):
            continue

        if i == 0:
            new_leading_lines = []
            found_comment = False
            for line in first_node.leading_lines:
                if line.comment is not None:
                    found_comment = True
                if found_comment:
                    new_leading_lines.append(line)
            block.nodes[0] = first_node.with_changes(leading_lines=new_leading_lines)
        else:
            if not first_node.leading_lines or first_node.leading_lines[0].comment is not None:
                block.nodes[0] = first_node.with_changes(
                    leading_lines=[cst.EmptyLine(indent=True, newline=cst.Newline())] + list(first_node.leading_lines)
                )

    new_body = []
    for b in sorted_blocks:
        new_body.extend(b.nodes)

    return module.with_changes(body=new_body).code


# ==============================================================================
# CLI & FILE PROCESSING
# ==============================================================================


def should_exclude(path: Path, exclude_patterns: List[str]) -> bool:
    path_str = str(path)
    for pattern in exclude_patterns:
        if fnmatch.fnmatch(path_str, pattern) or any(fnmatch.fnmatch(part, pattern) for part in path.parts):
            return True
    return False


def process_file(file_path: Path, config: Config) -> FileStatus:
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            source_code = f.read()
    except UnicodeDecodeError:
        if not config.quiet:
            print(f'Error ({file_path}): UnicodeDecodeError', file=sys.stderr)
        return FileStatus.ERROR

    try:
        formatted_code = format_code(source_code, config.alphabetical, config.sort_imports)
    except Exception as e:
        if not config.quiet:
            print(f'Error ({file_path}): {e}', file=sys.stderr)
        return FileStatus.ERROR

    if source_code != formatted_code:
        if config.diff:
            if not config.quiet:
                sys.stdout.writelines(
                    difflib.unified_diff(
                        source_code.splitlines(keepends=True),
                        formatted_code.splitlines(keepends=True),
                        fromfile=f'a/{file_path}',
                        tofile=f'b/{file_path}',
                    )
                )

        if not config.check and not config.diff:
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(formatted_code)
            if config.verbose and not config.quiet:
                print(f'Reordered {file_path}')

        return FileStatus.CHANGED

    if config.verbose and not config.quiet:
        print(f'Unchanged {file_path}')
    return FileStatus.UNCHANGED


def run_cli():
    parser = argparse.ArgumentParser(description='funkshway: Opinionated Python code reordering')
    parser.add_argument(
        '--alphabetical',
        '-a',
        action='store_true',
        help='Sort classes/functions alphabetically instead of topologically',
    )
    parser.add_argument('--sort-imports', '-i', action='store_true', help='Sort statements within import blocks')
    parser.add_argument('--exclude', type=str, help='Comma-separated glob patterns to exclude')
    parser.add_argument('--relaxed', action='store_true', help='Exit with zero even on parse errors')
    parser.add_argument('--check', action='store_true', help="Don't write files, exit 1 if changes needed")
    parser.add_argument('--diff', action='store_true', help="Don't write files, print unified diff")
    parser.add_argument('--verbose', '-v', action='store_true', help='Print details of unchanged files')
    parser.add_argument('--quiet', '-q', action='store_true', help='Suppress all non-error output')

    parser.add_argument('mode', nargs='?', choices=['.', 'files'], help='Mode. Omit to launch GUI.')
    parser.add_argument('paths', nargs='*', help="Files to format (use '-' for stdin)")

    args = parser.parse_args()

    toml_data = load_pyproject_toml()
    config = Config()
    config.alphabetical = args.alphabetical or toml_data.get('alphabetical', False)
    config.sort_imports = args.sort_imports or toml_data.get('sort_imports', False)
    config.check = args.check
    config.diff = args.diff
    config.relaxed = args.relaxed
    config.verbose = args.verbose
    config.quiet = args.quiet

    if args.exclude:
        config.exclude = args.exclude.split(',')
    elif 'exclude' in toml_data:
        config.exclude = toml_data['exclude']

    if not args.mode:
        if args.check or args.diff or args.relaxed or args.alphabetical or args.sort_imports:
            parser.error("Flags require a mode ('.' or 'files'). Omit all arguments to launch GUI.")
        run_gui()
        return

    if '-' in args.paths:
        try:
            source_code = sys.stdin.read()
            formatted_code = format_code(source_code, config.alphabetical, config.sort_imports)
        except Exception as e:
            if not config.quiet:
                print(f'Error processing stdin: {e}', file=sys.stderr)
            sys.exit(0 if config.relaxed else 1)

        if source_code != formatted_code:
            if config.diff and not config.quiet:
                sys.stdout.writelines(
                    difflib.unified_diff(
                        source_code.splitlines(keepends=True),
                        formatted_code.splitlines(keepends=True),
                        fromfile='a/stdin',
                        tofile='b/stdin',
                    )
                )
            elif not config.check and not config.diff:
                sys.stdout.write(formatted_code)
            sys.exit(1 if (config.check or config.diff) else 0)

        if not config.check and not config.diff:
            sys.stdout.write(source_code)
        sys.exit(0)

    files_to_format: List[Path] = []

    if args.mode == '.':
        if args.paths:
            parser.error('recursive mode does not accept additional arguments')
        for path in Path('.').rglob('*.py'):
            if should_exclude(path, config.exclude):
                continue
            files_to_format.append(path)

    elif args.mode == 'files':
        if not args.paths:
            parser.error("'files' mode requires at least one file")
        for path_str in args.paths:
            path = Path(path_str)
            if not path.is_file():
                if not config.quiet:
                    print(f'Error: Not a file: {path_str}', file=sys.stderr)
                sys.exit(1)
            files_to_format.append(path)

    has_changes = False
    has_errors = False

    with concurrent.futures.ProcessPoolExecutor() as executor:
        futures = {executor.submit(process_file, f, config): f for f in files_to_format}
        for future in concurrent.futures.as_completed(futures):
            status = future.result()
            if status == FileStatus.ERROR:
                has_errors = True
            elif status == FileStatus.CHANGED:
                has_changes = True

    if has_errors:
        sys.exit(0 if config.relaxed else 1)

    if has_changes and (config.check or config.diff):
        sys.exit(1)

    sys.exit(0)


# ==============================================================================
# GUI (TKINTER)
# ==============================================================================


def run_gui():
    root = tk.Tk()
    root.title('Funkshway - Code Reorder Debugger')
    root.geometry('950x700')

    font_style = ('Consolas', 11)

    var_alphabetical = tk.BooleanVar(value=False)
    var_sort_imports = tk.BooleanVar(value=False)

    top_frame = tk.Frame(root)
    top_frame.pack(fill=tk.X, padx=10, pady=5)

    lbl = tk.Label(top_frame, text='Paste your Python code below:', font=('Arial', 10, 'bold'))
    lbl.pack(side=tk.LEFT)

    chk_alpha = tk.Checkbutton(top_frame, text='Alphabetical Order', variable=var_alphabetical)
    chk_alpha.pack(side=tk.RIGHT, padx=5)

    chk_imports = tk.Checkbutton(top_frame, text='Sort Imports', variable=var_sort_imports)
    chk_imports.pack(side=tk.RIGHT, padx=5)

    text_area = tk.Text(root, wrap='none', font=font_style, undo=True)

    y_scroll = tk.Scrollbar(root, command=text_area.yview)
    y_scroll.pack(side=tk.RIGHT, fill=tk.Y)
    x_scroll = tk.Scrollbar(root, orient=tk.HORIZONTAL, command=text_area.xview)
    x_scroll.pack(side=tk.BOTTOM, fill=tk.X)

    text_area.configure(yscrollcommand=y_scroll.set, xscrollcommand=x_scroll.set)
    text_area.pack(expand=True, fill=tk.BOTH, padx=10, pady=5)

    status_bar = tk.Label(root, text='Ready', bd=1, relief=tk.SUNKEN, anchor=tk.W, font=('Arial', 10))
    status_bar.pack(side=tk.BOTTOM, fill=tk.X)

    def update_status(message: str, is_error: bool = False):
        color = '#c92a2a' if is_error else '#2b8a3e'
        status_bar.config(text=message, fg=color)

    def do_format(source_code: str, file_path: Optional[Path] = None):
        try:
            formatted = format_code(source_code, var_alphabetical.get(), var_sort_imports.get())
            changed = formatted != source_code

            text_area.delete('1.0', tk.END)
            text_area.insert(tk.END, formatted)

            if file_path:
                if changed:
                    with open(file_path, 'w', encoding='utf-8') as f:
                        f.write(formatted)
                    update_status(f"Success: File '{file_path.name}' formatted and saved.")
                else:
                    update_status(f"Unchanged: File '{file_path.name}' was already formatted.")
            else:
                if changed:
                    update_status('Success: Code successfully formatted!')
                else:
                    update_status('Unchanged: Code was already formatted.')

        except Exception as e:
            update_status('Error: Failed to format code.', is_error=True)
            messagebox.showerror('Parse Error', f'Failed to parse/format code:\n{e}')

    def paste_and_format():
        try:
            raw_text = root.clipboard_get().strip()
        except tk.TclError:
            update_status('Warning: Clipboard is empty or contains no text.', is_error=True)
            return

        if not raw_text:
            update_status('Warning: Clipboard text is empty.', is_error=True)
            return

        # Проверка, не является ли текст в буфере путем к файлу
        possible_path = None
        if len(raw_text.splitlines()) == 1:
            try:
                p = Path(raw_text)
                if p.is_file():
                    possible_path = p
            except OSError:
                pass

        if possible_path:
            try:
                with open(possible_path, 'r', encoding='utf-8') as f:
                    source_code = f.read()
                do_format(source_code, file_path=possible_path)
            except Exception as e:
                update_status(f"Error: Failed to read file '{possible_path.name}'.", is_error=True)
                messagebox.showerror('File Error', str(e))
        else:
            do_format(raw_text)

    def format_current():
        raw_text = text_area.get('1.0', tk.END).strip()
        if not raw_text:
            update_status('Warning: Text area is empty.', is_error=True)
            return
        do_format(raw_text)

    def copy_to_clipboard():
        text_content = text_area.get('1.0', tk.END).strip()
        if not text_content:
            update_status('Warning: Nothing to copy.', is_error=True)
            return
        root.clipboard_clear()
        root.clipboard_append(text_content)
        update_status('Success: Formatted code copied to clipboard!')

    btn_frame = tk.Frame(root)
    btn_frame.pack(fill=tk.X, padx=10, pady=10)

    btn_paste = tk.Button(btn_frame, text='📋 Paste & Format', command=paste_and_format, bg='#e0e0e0', height=2)
    btn_paste.pack(side=tk.LEFT, padx=5, expand=True, fill=tk.X)

    btn_format = tk.Button(btn_frame, text='✨ Format Current', command=format_current, bg='#d4edda', height=2)
    btn_format.pack(side=tk.LEFT, padx=5, expand=True, fill=tk.X)

    btn_copy = tk.Button(btn_frame, text='💾 Copy', command=copy_to_clipboard, bg='#d1ecf1', height=2)
    btn_copy.pack(side=tk.LEFT, padx=5, expand=True, fill=tk.X)

    btn_clear = tk.Button(
        btn_frame,
        text='🗑 Clear',
        command=lambda: [text_area.delete('1.0', tk.END), update_status('Ready')],
        bg='#f8d7da',
        height=2,
    )
    btn_clear.pack(side=tk.LEFT, padx=5, expand=True, fill=tk.X)

    root.mainloop()


if __name__ == '__main__':
    run_cli()