|
| 1 | +package main |
| 2 | + |
| 3 | +import ( |
| 4 | + "fmt" |
| 5 | + "io/ioutil" |
| 6 | + "os" |
| 7 | + |
| 8 | + sitter "github.com/smacker/go-tree-sitter" |
| 9 | + "github.com/smacker/go-tree-sitter/python" |
| 10 | + "github.com/smacker/go-tree-sitter/yaml" |
| 11 | + "github.com/smacker/go-tree-sitter/sql" |
| 12 | + "github.com/smacker/go-tree-sitter/bash" |
| 13 | + "github.com/smacker/go-tree-sitter/html" |
| 14 | + "github.com/smacker/go-tree-sitter/dockerfile" |
| 15 | +) |
| 16 | + |
| 17 | +func main() { |
| 18 | + // Check if both language and query string are provided as arguments |
| 19 | + if len(os.Args) < 3 { |
| 20 | + fmt.Println("Usage: ./program <language> <query_string>") |
| 21 | + return |
| 22 | + } |
| 23 | + language := os.Args[1] |
| 24 | + queryString := os.Args[2] |
| 25 | + |
| 26 | + // Create a parser |
| 27 | + parser := sitter.NewParser() |
| 28 | + |
| 29 | + // Set the language based on the first argument |
| 30 | + switch language { |
| 31 | + case "python": |
| 32 | + parser.SetLanguage(python.GetLanguage()) |
| 33 | + case "yaml": |
| 34 | + parser.SetLanguage(yaml.GetLanguage()) |
| 35 | + case "sql": |
| 36 | + parser.SetLanguage(sql.GetLanguage()) |
| 37 | + case "bash": |
| 38 | + parser.SetLanguage(bash.GetLanguage()) |
| 39 | + case "html": |
| 40 | + parser.SetLanguage(html.GetLanguage()) |
| 41 | + case "dockerfile": |
| 42 | + parser.SetLanguage(dockerfile.GetLanguage()) |
| 43 | + default: |
| 44 | + fmt.Printf("Unsupported language: %s\n", language) |
| 45 | + return |
| 46 | + } |
| 47 | + |
| 48 | + // Read source code from stdin |
| 49 | + sourceCode, err := ioutil.ReadAll(os.Stdin) |
| 50 | + if err != nil { |
| 51 | + fmt.Println("Error reading from stdin:", err) |
| 52 | + return |
| 53 | + } |
| 54 | + |
| 55 | + // Parse the source code |
| 56 | + tree := parser.Parse(nil, sourceCode) |
| 57 | + defer tree.Close() |
| 58 | + |
| 59 | + // Create a query |
| 60 | + query, err := sitter.NewQuery([]byte(queryString), python.GetLanguage()) |
| 61 | + if err != nil { |
| 62 | + fmt.Println("Error creating query:", err) |
| 63 | + return |
| 64 | + } |
| 65 | + defer query.Close() |
| 66 | + |
| 67 | + // Execute the query |
| 68 | + qc := sitter.NewQueryCursor() |
| 69 | + defer qc.Close() |
| 70 | + |
| 71 | + qc.Exec(query, tree.RootNode()) |
| 72 | + |
| 73 | + // Iterate over the query results |
| 74 | + for { |
| 75 | + match, ok := qc.NextMatch() |
| 76 | + if !ok { |
| 77 | + break |
| 78 | + } |
| 79 | + |
| 80 | + for _, capture := range match.Captures { |
| 81 | + captureName := query.CaptureNameForId(capture.Index) |
| 82 | + nodeText := capture.Node.Content(sourceCode) |
| 83 | + fmt.Printf("Capture: %s, Node: %s\n", captureName, nodeText) |
| 84 | + } |
| 85 | + } |
| 86 | +} |
0 commit comments