Coverage for .tox/py312/lib/python3.12/site-packages/pylint_sort_functions/auto_fix.py: 100%
176 statements
« prev ^ index » next coverage.py v7.10.2, created at 2025-08-07 04:45 +0200
1"""Auto-fix functionality for sorting functions and methods."""
3import shutil
4from dataclasses import dataclass
5from pathlib import Path
6from typing import List, Optional, Tuple
8import astroid # type: ignore[import-untyped]
9from astroid import nodes
11from pylint_sort_functions import utils
14@dataclass
15class FunctionSpan:
16 """Represents a function with its complete text span in the source file."""
18 node: nodes.FunctionDef
19 start_line: int
20 end_line: int
21 text: str # Complete source text from start_line to end_line (inclusive)
22 name: str
25@dataclass
26class AutoFixConfig:
27 """Configuration for the automatic function sorting tool.
29 Controls how the auto-fix feature behaves when reordering functions
30 and methods in Python source files.
32 Note: Comment preservation is always enabled as it's essential for
33 maintaining code intent and documentation during reorganization.
34 """
36 dry_run: bool = False # Show what would be changed without modifying files
37 backup: bool = True # Create .bak files before making changes
38 ignore_decorators: Optional[List[str]] = (
39 None # Decorator patterns to exclude from sorting
40 )
43# Note: This class intentionally has only one public method as it encapsulates
44# the configuration state and provides a clean interface for file processing.
45class FunctionSorter: # pylint: disable=too-few-public-methods
46 """Main class for auto-fixing function sorting.
48 This class provides the core functionality for automatically reordering
49 functions and methods in Python source files to comply with sorting rules.
51 Usage:
52 Used by the CLI tool (cli.py) and can be used programmatically:
54 config = AutoFixConfig(dry_run=True, backup=True)
55 sorter = FunctionSorter(config)
56 was_modified = sorter.sort_file(Path("my_file.py"))
57 """
59 def __init__(self, config: AutoFixConfig):
60 """Initialize the function sorter.
62 :param config: Configuration for auto-fix behavior
63 :type config: AutoFixConfig
64 """
65 self.config = config
66 if self.config.ignore_decorators is None:
67 self.config.ignore_decorators = []
69 def sort_file(self, file_path: Path) -> bool:
70 """Auto-sort functions in a Python file.
72 :param file_path: Path to the Python file to sort
73 :type file_path: Path
74 :returns: True if file was modified, False otherwise
75 :rtype: bool
76 """
77 try:
78 # Read the original file
79 original_content = file_path.read_text(encoding="utf-8")
81 # Check if file needs sorting
82 if not self._file_needs_sorting(original_content):
83 return False
85 # Extract and sort functions
86 new_content = self._sort_functions_in_content(original_content)
88 if new_content == original_content: # pragma: no cover
89 return False
91 if self.config.dry_run:
92 print(f"Would modify: {file_path}")
93 return True
95 # Create backup if requested
96 if self.config.backup:
97 backup_path = file_path.with_suffix(f"{file_path.suffix}.bak")
98 shutil.copy2(file_path, backup_path)
100 # Write the sorted content
101 file_path.write_text(new_content, encoding="utf-8")
102 return True
104 # Broad exception catch ensures tool never crashes when modifying user files
105 except (
106 Exception
107 ) as e: # pragma: no cover # pylint: disable=broad-exception-caught
108 print(f"Error processing {file_path}: {e}")
109 return False
111 def _extract_function_spans(
112 self, functions: List[nodes.FunctionDef], lines: List[str]
113 ) -> List[FunctionSpan]:
114 """Extract function text spans from the source.
116 :param functions: List of function nodes
117 :type functions: List[nodes.FunctionDef]
118 :param lines: Source file lines
119 :type lines: List[str]
120 :returns: List of function spans with text
121 :rtype: List[FunctionSpan]
122 """
123 spans = []
125 # First pass: determine where each function (including comments) starts
126 function_boundaries = []
127 for func in functions:
128 start_line = func.lineno - 1 # Convert to 0-based indexing
130 # Include decorators in the span
131 actual_start = start_line
132 if hasattr(func, "decorators") and func.decorators:
133 actual_start = func.decorators.lineno - 1
135 # Include comments above the function/decorators
136 comment_start = self._find_comments_above_function(lines, actual_start)
137 function_boundaries.append((func, comment_start))
139 # Second pass: create spans using the boundaries
140 for i, (func, comment_start) in enumerate(function_boundaries):
141 # Find the end line (start of next function or end of file)
142 if i + 1 < len(function_boundaries):
143 # End where the next function's comments start
144 end_line = function_boundaries[i + 1][1]
145 else:
146 # Last function, use end of file
147 end_line = len(lines)
149 # Extract the text including comments
150 text = "".join(lines[comment_start:end_line])
152 spans.append(
153 FunctionSpan(
154 node=func,
155 start_line=comment_start,
156 end_line=end_line,
157 text=text,
158 name=func.name,
159 )
160 )
162 return spans
164 def _extract_method_spans(
165 self,
166 methods: List[nodes.FunctionDef],
167 lines: List[str],
168 class_node: nodes.ClassDef,
169 ) -> List[FunctionSpan]:
170 """Extract method text spans from a class.
172 :param methods: List of method nodes from the class
173 :type methods: List[nodes.FunctionDef]
174 :param lines: Source file lines
175 :type lines: List[str]
176 :param class_node: The class containing these methods
177 :type class_node: nodes.ClassDef
178 :returns: List of method spans with text
179 :rtype: List[FunctionSpan]
180 """
181 spans = []
183 # First pass: determine where each method (including comments) starts
184 method_boundaries = []
185 for method in methods:
186 start_line = method.lineno - 1 # Convert to 0-based indexing
188 # Include decorators in the span
189 actual_start = start_line
190 if hasattr(method, "decorators") and method.decorators:
191 actual_start = method.decorators.lineno - 1
193 # Include comments above the method/decorators
194 comment_start = self._find_comments_above_function(lines, actual_start)
195 method_boundaries.append((method, comment_start))
197 # Second pass: create spans using the boundaries
198 for i, (method, comment_start) in enumerate(method_boundaries):
199 # Find the end line (start of next method or end of class)
200 if i + 1 < len(method_boundaries):
201 # End where the next method's comments start
202 end_line = method_boundaries[i + 1][1]
203 else:
204 # Last method in class, find end of class
205 end_line = (
206 class_node.end_lineno
207 if hasattr(class_node, "end_lineno")
208 else len(lines)
209 )
211 # Extract the text including comments
212 text = "".join(lines[comment_start:end_line])
214 spans.append(
215 FunctionSpan(
216 node=method,
217 start_line=comment_start,
218 end_line=end_line,
219 text=text,
220 name=method.name,
221 )
222 )
224 return spans
226 def _file_needs_sorting(self, content: str) -> bool:
227 """Check if a file needs function sorting.
229 :param content: File content as string
230 :type content: str
231 :returns: True if file needs sorting
232 :rtype: bool
233 """
234 try:
235 # Parse with astroid for consistency with the checker
236 module = astroid.parse(content)
238 # Check module-level functions
239 functions = utils.get_functions_from_node(module)
240 if functions and not utils.are_functions_sorted_with_exclusions(
241 functions, self.config.ignore_decorators
242 ):
243 return True
245 # Check class methods
246 for node in module.body:
247 if isinstance(node, nodes.ClassDef):
248 methods = utils.get_methods_from_class(node)
249 if methods and not utils.are_methods_sorted_with_exclusions(
250 methods, self.config.ignore_decorators
251 ):
252 return True
254 return False
256 except Exception: # pylint: disable=broad-exception-caught
257 return False
259 def _find_comments_above_function(
260 self, lines: List[str], function_start_line: int
261 ) -> int:
262 """Find comments that belong to a function and return the start line.
264 Scans backwards from the function definition to find associated comments.
266 :param lines: Source file lines
267 :type lines: List[str]
268 :param function_start_line: The line where the function starts (0-based)
269 :type function_start_line: int
270 :returns: The line number where comments start, or function_start_line
271 :rtype: int
272 """
273 comment_start_line = function_start_line
275 # Scan backwards from the function start to find comments
276 current_line = function_start_line - 1
278 while current_line >= 0:
279 line = lines[current_line].strip()
281 # If we find a comment line, this could be part of the function's comments
282 if line.startswith("#"):
283 comment_start_line = current_line
284 current_line -= 1
285 continue
287 # If we find an empty line, continue scanning (comments might be separated)
288 if line == "":
289 current_line -= 1
290 continue
292 # If we find any other content, stop scanning
293 break
295 return comment_start_line
297 def _reconstruct_class_with_sorted_methods(
298 self,
299 content: str,
300 original_spans: List[FunctionSpan],
301 sorted_spans: List[FunctionSpan],
302 ) -> str:
303 """Reconstruct class content with sorted methods.
305 :param content: Original file content
306 :type content: str
307 :param original_spans: Original method spans in order of appearance
308 :type original_spans: List[FunctionSpan]
309 :param sorted_spans: Method spans in sorted order
310 :type sorted_spans: List[FunctionSpan]
311 :returns: Reconstructed content with sorted methods
312 :rtype: str
313 """
314 if not original_spans: # pragma: no cover
315 return content
317 # Find the region that contains all methods within the class
318 first_method_start = min(span.start_line for span in original_spans)
319 last_method_end = max(span.end_line for span in original_spans)
321 # Split content into lines for manipulation
322 content_lines = content.splitlines(keepends=True)
324 # Build new content
325 new_lines = []
327 # Add everything before the first method
328 new_lines.extend(content_lines[:first_method_start])
330 # Add sorted methods
331 for span in sorted_spans:
332 # Method text already includes proper spacing
333 new_lines.append(span.text)
335 # Add everything after the last method
336 if last_method_end < len(content_lines):
337 new_lines.extend(content_lines[last_method_end:])
339 return "".join(new_lines)
341 def _reconstruct_content_with_sorted_functions(
342 self,
343 original_content: str,
344 original_spans: List[FunctionSpan],
345 sorted_spans: List[FunctionSpan],
346 ) -> str:
347 """Reconstruct file content with sorted functions.
349 Strategy:
350 1. Preserve everything before the first function (imports, module docstrings)
351 2. Replace the function block with sorted functions
352 3. Preserve everything after the last function
353 4. Add blank lines between functions if not already present
355 This approach ensures non-function content (imports, constants, etc.)
356 remains in its original position while only reordering functions.
358 :param original_content: Original file content
359 :type original_content: str
360 :param original_spans: Original function spans in order of appearance
361 :type original_spans: List[FunctionSpan]
362 :param sorted_spans: Function spans in sorted order
363 :type sorted_spans: List[FunctionSpan]
364 :returns: Reconstructed content with sorted functions
365 :rtype: str
366 """
367 if not original_spans: # pragma: no cover
368 return original_content
370 lines = original_content.splitlines(keepends=True)
372 # Find the region that contains all functions
373 first_func_start = min(span.start_line for span in original_spans)
374 last_func_end = max(span.end_line for span in original_spans)
376 # Build new content
377 new_lines = []
379 # Add everything before the first function
380 new_lines.extend(lines[:first_func_start])
382 # Add sorted functions
383 for i, span in enumerate(sorted_spans):
384 if i > 0:
385 # Add blank line between functions if not already present
386 if not span.text.startswith("\n"):
387 new_lines.append("\n")
388 new_lines.append(span.text)
390 # Add everything after the last function
391 if last_func_end < len(lines): # pragma: no cover
392 new_lines.extend(lines[last_func_end:])
394 return "".join(new_lines)
396 def _sort_class_methods(
397 self, content: str, module: nodes.Module, lines: List[str]
398 ) -> str:
399 """Sort methods within classes.
401 :param content: File content
402 :type content: str
403 :param module: Parsed module
404 :type module: nodes.Module
405 :param lines: Content split into lines
406 :type lines: List[str]
407 :returns: Content with sorted class methods
408 :rtype: str
409 """
410 # Find all classes that need method sorting
411 classes_to_sort = []
412 for node in module.body:
413 if isinstance(node, nodes.ClassDef):
414 methods = utils.get_methods_from_class(node)
415 if methods and not utils.are_methods_sorted_with_exclusions(
416 methods, self.config.ignore_decorators
417 ):
418 classes_to_sort.append((node, methods))
420 if not classes_to_sort:
421 return content
423 # Sort each class's methods
424 for class_node, methods in classes_to_sort:
425 # Extract method spans for this class
426 method_spans = self._extract_method_spans(methods, lines, class_node)
428 # Sort the method spans
429 sorted_spans = self._sort_function_spans(method_spans)
431 # Reconstruct the class content with sorted methods
432 content = self._reconstruct_class_with_sorted_methods(
433 content, method_spans, sorted_spans
434 )
436 return content
438 def _sort_function_spans(self, spans: List[FunctionSpan]) -> List[FunctionSpan]:
439 """Sort function spans according to the plugin rules.
441 :param spans: List of function spans to sort
442 :type spans: List[FunctionSpan]
443 :returns: Sorted list of function spans
444 :rtype: List[FunctionSpan]
445 """
446 # Separate functions based on exclusions and visibility
447 excluded = []
448 sortable_public = []
449 sortable_private = []
451 for span in spans:
452 if utils.function_has_excluded_decorator(
453 span.node, self.config.ignore_decorators or []
454 ):
455 excluded.append(span)
456 elif utils.is_private_function(span.node):
457 sortable_private.append(span)
458 else:
459 sortable_public.append(span)
461 # Sort the sortable functions alphabetically
462 sortable_public.sort(key=lambda s: s.name)
463 sortable_private.sort(key=lambda s: s.name)
465 # Reconstruct the order: sortable public + sortable private + excluded
466 # For now, use a simple approach: public sorted + private sorted + excluded
467 # Future enhancement: Preserve relative positions of excluded functions
468 return sortable_public + sortable_private + excluded
470 def _sort_functions_in_content(self, content: str) -> str:
471 """Sort functions in file content and return new content.
473 :param content: Original file content
474 :type content: str
475 :returns: Content with sorted functions
476 :rtype: str
477 """
478 try:
479 module = astroid.parse(content)
480 lines = content.splitlines(keepends=True)
482 # Process module-level functions
483 content = self._sort_module_functions(content, module, lines)
485 # Process class methods
486 content = self._sort_class_methods(content, module, lines)
488 return content
490 except (
491 Exception
492 ) as e: # pragma: no cover # pylint: disable=broad-exception-caught
493 print(f"Error sorting content: {e}")
494 return content
496 def _sort_module_functions(
497 self, content: str, module: nodes.Module, lines: List[str]
498 ) -> str:
499 """Sort module-level functions.
501 :param content: File content
502 :type content: str
503 :param module: Parsed module
504 :type module: nodes.Module
505 :param lines: Content split into lines
506 :type lines: List[str]
507 :returns: Content with sorted module functions
508 :rtype: str
509 """
510 functions = utils.get_functions_from_node(module)
511 if not functions: # pragma: no cover
512 return content
514 # Check if sorting is needed
515 if utils.are_functions_sorted_with_exclusions( # pragma: no cover
516 functions, self.config.ignore_decorators
517 ):
518 return content
520 # Extract function spans
521 function_spans = self._extract_function_spans(functions, lines)
523 # Sort the functions
524 sorted_spans = self._sort_function_spans(function_spans)
526 # Reconstruct content
527 return self._reconstruct_content_with_sorted_functions(
528 content, function_spans, sorted_spans
529 )
532# Public API function for sorting a single file
533def sort_python_file(file_path: Path, config: AutoFixConfig) -> bool: # pylint: disable=function-should-be-private
534 """Sort functions in a Python file.
536 :param file_path: Path to the Python file
537 :type file_path: Path
538 :param config: Auto-fix configuration
539 :type config: AutoFixConfig
540 :returns: True if file was modified
541 :rtype: bool
542 """
543 return _sort_python_file(file_path, config)
546def sort_python_files(file_paths: List[Path], config: AutoFixConfig) -> Tuple[int, int]:
547 """Sort functions in multiple Python files.
549 :param file_paths: List of Python file paths
550 :type file_paths: List[Path]
551 :param config: Auto-fix configuration
552 :type config: AutoFixConfig
553 :returns: Tuple of (files_processed, files_modified)
554 :rtype: Tuple[int, int]
555 """
556 files_processed = 0
557 files_modified = 0
559 for file_path in file_paths:
560 if file_path.suffix == ".py":
561 files_processed += 1
562 if _sort_python_file(file_path, config):
563 files_modified += 1
565 return files_processed, files_modified
568# Private functions
571def _sort_python_file(file_path: Path, config: AutoFixConfig) -> bool:
572 """Sort functions in a Python file (private implementation).
574 :param file_path: Path to the Python file
575 :type file_path: Path
576 :param config: Auto-fix configuration
577 :type config: AutoFixConfig
578 :returns: True if file was modified
579 :rtype: bool
580 """
581 sorter = FunctionSorter(config)
582 return sorter.sort_file(file_path)