learningcircuit--local-deep-research/.pre-commit-hooks/check-session-context-manager.py

#!/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())