add yaml rule mcp

Author: HerringtonDarkholme

main.py

@@ -13,10 +13,17 @@ def find_code(
     pattern: str = Field(description="The ast-grep pattern to search for"),
     language: str = Field(description="The language of the query", default=""),
 ) -> List[dict[str, Any]]:
-    """Find code in a project folder that matches the given ast-grep pattern.
-    """
+    """Find code in a project folder that matches the given ast-grep pattern"""
     return run_ast_grep_command(pattern, project_folder, language)
 
+@mcp.tool()
+def find_code_by_rule(
+    project_folder: str = Field(description="The path to the project folder"),
+    yaml: str = Field(description="The ast-grep pattern to search for"),
+    ) -> List[dict[str, Any]]:
+    """Find code using ast-grep's YAML rule in a project folder"""
+    return run_ast_grep_yaml(yaml, project_folder)
+
 def run_ast_grep_command(pattern: str, project_folder: str, language: Optional[str]) -> List[dict[str, Any]]:
     try:
         args = ["ast-grep", "--pattern", pattern, "--json", project_folder]
@@ -38,5 +45,24 @@ def run_ast_grep_command(pattern: str, project_folder: str, language: Optional[s
         print("Command not found")
         return []
 
+def run_ast_grep_yaml(yaml: str, project_folder: str) -> List[dict[str, Any]]:
+    try:
+        args = ["ast-grep", "scan","--inline-rules", yaml, "--json", project_folder]
+        # Run command and capture output
+        result = subprocess.run(
+            args,
+            capture_output=True,
+            text=True,
+            check=True  # Raises CalledProcessError if return code is non-zero
+        )
+        return json.loads(result.stdout)
+    except subprocess.CalledProcessError as e:
+        print(f"Command failed with return code {e.returncode}")
+        print("Error output:", e.stderr)
+        return e.stderr
+    except FileNotFoundError:
+        print("Command not found")
+        return []
+
 if __name__ == "__main__":
     mcp.run(transport = "stdio")