#!/usr/bin/env python3 """ Pre-commit hook to detect try/finally session patterns and suggest context managers. This hook checks for SQLAlchemy session management patterns that use try/finally blocks and suggests replacing them with context managers for better resource management and cleaner code. """ import ast import sys from pathlib import Path from typing import List, Tuple class SessionPatternChecker(ast.NodeVisitor): """AST visitor to detect try/finally session patterns.""" def __init__(self, filename: str): self.filename = filename self.issues: List[Tuple[int, str]] = [] self.functions_and_methods = [] def visit_FunctionDef(self, node: ast.FunctionDef) -> None: """Visit function definitions to check for session patterns.""" self._check_function_for_pattern(node) self.generic_visit(node) def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: """Visit async function definitions to check for session patterns.""" self._check_function_for_pattern(node) self.generic_visit(node) def _check_function_for_pattern(self, func_node) -> None: """Check a function body for try/finally session patterns.""" # Look for session = Session() followed by try/finally for i, stmt in enumerate(func_node.body): # Check if this is a session assignment if isinstance(stmt, ast.Assign): session_var = self._get_session_var_from_assign(stmt) if session_var: # Look for a try/finally block that follows for next_stmt in func_node.body[ i + 1 : i + 3 ]: # Check next 2 statements if ( isinstance(next_stmt, ast.Try) and next_stmt.finalbody ): # Check if finally has session.close() if self._has_session_close_in_finally( next_stmt.finalbody, session_var ): self.issues.append( ( stmt.lineno, f"Found try/finally session pattern. Consider using 'with self.Session() as {session_var}:' instead", ) ) break def _get_session_var_from_assign(self, assign_node: ast.Assign) -> str: """Check if an assignment is creating a session and return the variable name.""" if isinstance(assign_node.value, ast.Call) and self._is_session_call( assign_node.value ): if assign_node.targets and isinstance( assign_node.targets[0], ast.Name ): return assign_node.targets[0].id return None def _is_session_call(self, call_node: ast.Call) -> bool: """Check if a call node is creating a SQLAlchemy session.""" # Check for self.Session() pattern if isinstance(call_node.func, ast.Attribute): if call_node.func.attr in ( "Session", "get_session", "create_session", ): return True # Check for Session() pattern elif isinstance(call_node.func, ast.Name): if call_node.func.id in ( "Session", "get_session", "create_session", ): return True return False def _has_session_close_in_finally( self, finalbody: List[ast.stmt], session_var: str ) -> bool: """Check if finally block contains session.close().""" for stmt in finalbody: if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call): # Check for session.close() pattern if ( isinstance(stmt.value.func, ast.Attribute) and stmt.value.func.attr == "close" ): # Check if it's our session variable if ( isinstance(stmt.value.func.value, ast.Name) and stmt.value.func.value.id == session_var ): return True return False def check_file(filepath: Path) -> List[Tuple[str, int, str]]: """Check a single Python file for try/finally session patterns.""" issues = [] try: content = filepath.read_text(encoding="utf-8") tree = ast.parse(content, filename=str(filepath)) checker = SessionPatternChecker(str(filepath)) checker.visit(tree) for line_no, message in checker.issues: issues.append((str(filepath), line_no, message)) except SyntaxError as e: # Skip files with syntax errors print(f"Syntax error in {filepath}: {e}", file=sys.stderr) except Exception as e: print(f"Error checking {filepath}: {e}", file=sys.stderr) return issues def main(): """Main entry point for the pre-commit hook.""" # Get list of files to check from command line arguments files_to_check = sys.argv[1:] if len(sys.argv) > 1 else [] if not files_to_check: print("No files to check") return 0 all_issues = [] for filepath_str in files_to_check: filepath = Path(filepath_str) # Skip non-Python files if not filepath.suffix == ".py": continue # Skip test files and migration files if "test" in filepath.parts or "migration" in filepath.parts: continue issues = check_file(filepath) all_issues.extend(issues) # Report issues if all_issues: print( "\n❌ Found try/finally session patterns that should use context managers:\n" ) for filepath, line_no, message in all_issues: print(f" {filepath}:{line_no}: {message}") print("\n💡 Tip: Replace try/finally blocks with context managers:") print(" Before:") print(" session = self.Session()") print(" try:") print(" # operations") print(" session.commit()") print(" finally:") print(" session.close()") print("\n After:") print(" with self.Session() as session:") print(" # operations") print(" session.commit()") print("\n") return 1 return 0 if __name__ == "__main__": sys.exit(main())