Add pre-commit hooks and scripts for async method checks in PGMQueue

jason810496 · jason810496 · commit d6fc36e8a7ba · 2026-01-05T15:45:19.000+08:00
- Introduced a pre-commit hook to check for missing async methods in PGMQueue.
- Added scripts to identify and generate missing async methods.
- Created utility functions for AST manipulation and method transformation.
- Established configuration for project paths and console output.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -7,4 +7,10 @@ repos:
       - id: ruff
         args: [ --fix ]
       # Run the formatter.
-      - id: ruff-format
+      - id: ruff-format
+  - repo: local
+    hooks:
+      - id: check-sync-async-method-for-queue
+        name: Check sync/async method for queue
+        entry: ./scripts/ci/pre_commit/check_sync_async_method_for_queue.py
+        language: python
diff --git a/scripts/ci/pre_commit/check_sync_async_method_for_queue.py b/scripts/ci/pre_commit/check_sync_async_method_for_queue.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+# /// script
+# requires-python = ">=3.10,<3.11"
+# dependencies = [
+#   "rich>=13.6.0",
+# ]
+# ///
+"""
+Script to check for missing async methods in PGMQueue for per-commit.
+
+For each public sync method (not starting with _), checks if there's a corresponding
+async method with the same name plus '_async' suffix.
+"""
+
+import ast
+import sys
+from pathlib import Path
+
+
+sys.path.insert(0, str(Path(__name__).parent.parent.joinpath("scripts").resolve()))
+
+from scripts_utils.config import QUEUE_FILE  # noqa: E402
+from scripts_utils.console import console  # noqa: E402
+from scripts_utils.common_ast import parse_methods_info_from_target_class  # noqa: E402
+
+
+def main():
+    """Main function."""
+
+    module_tree = ast.parse(source=QUEUE_FILE.read_text(), filename=QUEUE_FILE)
+    _, missing_async = parse_methods_info_from_target_class(
+        module_tree, target_class="PGMQueue"
+    )
+
+    if not missing_async:
+        console.print(
+            "[bold green]SUCCESS:[/bold green] All public methods have corresponding async versions!"
+        )
+        sys.exit(0)
+
+    # log all the missing async methods
+    console.print()
+    console.print(
+        f"[bold yellow]WARNING:[/bold yellow] Found {len(missing_async)} missing async methods:",
+        style="bold",
+    )
+    for method in missing_async:
+        console.print(f"  [yellow]-[/yellow] {method}_async")
+    console.print()
+
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/compelete_missing_async_methods.py b/scripts/compelete_missing_async_methods.py
@@ -0,0 +1,96 @@
+#!/usr/bin/env python
+# /// script
+# requires-python = ">=3.10,<3.11"
+# dependencies = [
+#   "rich>=13.6.0",
+# ]
+# ///
+"""
+Script to check for missing async methods in PGMQueue class and generate them.
+
+For each public sync method (not starting with _), checks if there's a corresponding
+async method with the same name plus '_async' suffix. If missing, generates it.
+"""
+
+import ast
+import sys
+from pathlib import Path
+import contextlib
+
+import tempfile
+
+
+from scripts_utils.config import QUEUE_FILE
+from scripts_utils.console import console
+from scripts_utils.common_ast import (
+    parse_methods_info_from_target_class,
+    fill_missing_methods_to_class,
+)
+from scripts_utils.formatting import format_file, compare_file
+from scripts_utils.queue_ast import get_async_methods_to_add
+
+
+def main():
+    """Main function."""
+
+    module_tree = ast.parse(source=QUEUE_FILE.read_text(), filename=QUEUE_FILE)
+    sync_methods, missing_async = parse_methods_info_from_target_class(
+        module_tree, target_class="PGMQueue"
+    )
+
+    if not missing_async:
+        console.print(
+            "[bold green]SUCCESS:[/bold green] All public methods have corresponding async versions!"
+        )
+        sys.exit(0)
+
+    # log all the missing async methods
+    console.print()
+    console.print(
+        f"[bold yellow]WARNING:[/bold yellow] Found {len(missing_async)} missing async methods:",
+        style="bold",
+    )
+    for method in missing_async:
+        console.print(f"  [yellow]-[/yellow] {method}_async")
+    console.print()
+
+    # create missing async method from
+    async_methods_to_add = get_async_methods_to_add(sync_methods, missing_async)
+    # insert back to class
+    fill_missing_methods_to_class(module_tree, "PGMQueue", async_methods_to_add)
+    module_tree = ast.fix_missing_locations(module_tree)
+
+    # write back to tmp file for comparison
+    tmp_file = ""
+    with tempfile.NamedTemporaryFile(mode="w+t", delete=False, suffix=".py") as f:
+        f.write(ast.unparse(module_tree))
+        f.flush()
+        tmp_file = f.name
+        console.log(f"Complete missing async methods at {tmp_file}")
+
+    if tmp_file:
+        max_formatting = 3
+        for _ in range(max_formatting):
+            if format_file(tmp_file):
+                break
+
+    _, missing_async_for_tmp = parse_methods_info_from_target_class(
+        ast.parse(Path(tmp_file).read_text()), "PGMQueue"
+    )
+
+    if missing_async_for_tmp:
+        console.log(
+            f"[error]Still get async methods to add after generating missing async methods in {tmp_file}: {missing_async_for_tmp}[/]"
+        )
+    else:
+        console.log("[success]All missing async methods are generated[/]")
+
+    # compare existed queue.py and tmp.py
+    with contextlib.suppress(Exception):
+        compare_file(QUEUE_FILE, tmp_file)
+
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/scripts_utils/__init__.py b/scripts/scripts_utils/__init__.py
diff --git a/scripts/scripts_utils/common_ast.py b/scripts/scripts_utils/common_ast.py
@@ -0,0 +1,124 @@
+from typing import List, Dict, Tuple, Literal
+
+import ast
+
+
+class MethodInfo:
+    """Information about a method."""
+
+    def __init__(self, name: str, node: ast.FunctionDef):
+        self.name = name
+        self.node = node
+        self.is_target = not name.startswith(
+            "_"
+        )  # all the public method is our target method for further processing
+        self.is_async = name.endswith("_async")
+        self.base_name = name[:-6] if self.is_async else name
+
+
+class ParseTargetClassFunctionsVisitor(ast.NodeVisitor):
+    """AST visitor to parse functions out of target class name for given module tree"""
+
+    def __init__(self, class_name: str):
+        self.class_name = class_name
+        self.methods: List[MethodInfo] = []
+        self.is_cur_node_in_target_class = False
+
+    def visit_ClassDef(self, node: ast.ClassDef):
+        if node.name == self.class_name:
+            self.is_cur_node_in_target_class = True
+            self.generic_visit(node)
+            self.is_cur_node_in_target_class = False
+        else:
+            self.generic_visit(node)
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        if self.is_cur_node_in_target_class:
+            # add all the method to the methods
+            self.methods.append(MethodInfo(node.name, node))
+        self.generic_visit(node)
+
+    def visit_AsyncFunctionDef(self, node):
+        if self.is_cur_node_in_target_class:
+            # add all the method to the methods
+            self.methods.append(MethodInfo(node.name, node))
+        self.generic_visit(node)
+
+
+class FillMissingMethodsToClass(ast.NodeTransformer):
+    """AST Transformer to fill missing async_methods back to target class"""
+
+    def __init__(self, class_name: str, to_add_async_methods: Dict[str, MethodInfo]):
+        self.class_name = class_name
+        self.to_add_async_methods = to_add_async_methods
+
+    def visit_ClassDef(self, node: ast.ClassDef):
+        if node.name == self.class_name:
+            for sync_func_name, async_func_node in self.to_add_async_methods.items():
+                idx = next(
+                    (
+                        i
+                        for i, stmt in enumerate(node.body)
+                        if isinstance(stmt, ast.FunctionDef)
+                        and stmt.name == sync_func_name
+                    ),
+                    -1,
+                )
+
+                if idx != -1:
+                    node.body.insert(idx + 1, async_func_node.node)
+
+        return self.generic_visit(node)
+
+
+def parse_methods_info_from_target_class(
+    module_tree: ast.Module, target_class: Literal["PGMQueue", "PGMQOperation"]
+) -> Tuple[List[MethodInfo], set[str]]:
+    """
+    Parse methods of target class from give module AST Tree
+
+    Args:
+        module_tree: ast.Module
+        target_class: either "PGMQueue" or "PGMQOperation" str
+
+    Returns:
+        Tuple of sync_methods, missing_async_set
+    """
+
+    analyzer = ParseTargetClassFunctionsVisitor(target_class)
+    analyzer.visit(module_tree)
+
+    # Categorize methods
+    # We use sync methods as source of truth
+    async_methods_set = set()
+    missing_async_set = set()
+
+    for method_info in analyzer.methods:
+        # skip non target methods
+        if not method_info.is_target:
+            continue
+
+        if method_info.is_async:
+            async_methods_set.add(method_info.base_name)
+
+    # Find missing async methods and generate class with interleaved methods
+    for method_info in analyzer.methods:
+        # skip non target methods
+        if not method_info.is_target:
+            continue
+
+        if method_info.base_name not in async_methods_set:
+            missing_async_set.add(method_info.base_name)
+
+    return analyzer.methods, missing_async_set
+
+
+def fill_missing_methods_to_class(
+    module_tree: ast.Module,
+    target_class: Literal["PGMQueue", "PGMQOperation"],
+    to_add_async_methods: Dict[str, MethodInfo],
+):
+    transformer = FillMissingMethodsToClass(
+        class_name=target_class, to_add_async_methods=to_add_async_methods
+    )
+    transformer.visit(module_tree)
diff --git a/scripts/scripts_utils/config.py b/scripts/scripts_utils/config.py
@@ -0,0 +1,5 @@
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).parent.parent.parent
+SOURCE_PATH = PROJECT_ROOT / "pgmq_sqlalchemy"
+QUEUE_FILE = SOURCE_PATH / "queue.py"
diff --git a/scripts/scripts_utils/console.py b/scripts/scripts_utils/console.py
@@ -0,0 +1,18 @@
+from rich.console import Console
+from rich.theme import Theme
+
+
+console = Console(
+    force_terminal=True,
+    color_system="standard",
+    theme=Theme(
+        {
+            "success": "green",
+            "info": "bright_blue",
+            "warning": "bright_yellow",
+            "error": "red",
+            "special": "magenta",
+        }
+    ),
+    width=202,
+)
diff --git a/scripts/scripts_utils/formatting.py b/scripts/scripts_utils/formatting.py
@@ -0,0 +1,24 @@
+import subprocess
+import sys
+from pathlib import Path
+
+
+sys.path.insert(0, str(Path(__name__).parent.parent.joinpath("scripts").resolve()))
+
+
+def format_file(file_path: str) -> bool:
+    try:
+        ruff_stdout = subprocess.check_output(["ruff", "format", file_path]).decode()
+    except Exception as e:
+        raise e
+
+    return "unchanged" in ruff_stdout
+
+
+def compare_file(existed_file: str, new_file: str):
+    try:
+        subprocess.check_call(
+            ["git", "difftool", "--tool=vimdiff", "--no-index", existed_file, new_file]
+        )
+    except Exception as e:
+        raise e
diff --git a/scripts/scripts_utils/queue_ast.py b/scripts/scripts_utils/queue_ast.py
