1package syntax_test
2
3import (
4	"bytes"
5	"fmt"
6	"log"
7	"reflect"
8	"strings"
9	"testing"
10
11	"go.starlark.net/syntax"
12)
13
14func TestWalk(t *testing.T) {
15	const src = `
16for x in y:
17  if x:
18    pass
19  else:
20    f([2*x for x in "abc"])
21`
22	// TODO(adonovan): test that it finds all syntax.Nodes
23	// (compare against a reflect-based implementation).
24	// TODO(adonovan): test that the result of f is used to prune
25	// the descent.
26	f, err := syntax.Parse("hello.go", src, 0)
27	if err != nil {
28		t.Fatal(err)
29	}
30
31	var buf bytes.Buffer
32	var depth int
33	syntax.Walk(f, func(n syntax.Node) bool {
34		if n == nil {
35			depth--
36			return true
37		}
38		fmt.Fprintf(&buf, "%s%s\n",
39			strings.Repeat("  ", depth),
40			strings.TrimPrefix(reflect.TypeOf(n).String(), "*syntax."))
41		depth++
42		return true
43	})
44	got := buf.String()
45	want := `
46File
47  ForStmt
48    Ident
49    Ident
50    IfStmt
51      Ident
52      BranchStmt
53      ExprStmt
54        CallExpr
55          Ident
56          Comprehension
57            BinaryExpr
58              Literal
59              Ident
60            ForClause
61              Ident
62              Literal`
63	got = strings.TrimSpace(got)
64	want = strings.TrimSpace(want)
65	if got != want {
66		t.Errorf("got %s, want %s", got, want)
67	}
68}
69
70// ExampleWalk demonstrates the use of Walk to
71// enumerate the identifiers in a Starlark source file
72// containing a nonsense program with varied grammar.
73func ExampleWalk() {
74	const src = `
75load("library", "a")
76
77def b(c, *, d=e):
78    f += {g: h}
79    i = -(j)
80    return k.l[m + n]
81
82for o in [p for q, r in s if t]:
83    u(lambda: v, w[x:y:z])
84`
85	f, err := syntax.Parse("hello.star", src, 0)
86	if err != nil {
87		log.Fatal(err)
88	}
89
90	var idents []string
91	syntax.Walk(f, func(n syntax.Node) bool {
92		if id, ok := n.(*syntax.Ident); ok {
93			idents = append(idents, id.Name)
94		}
95		return true
96	})
97	fmt.Println(strings.Join(idents, " "))
98
99	// The identifer 'a' appears in both LoadStmt.From[0] and LoadStmt.To[0].
100
101	// Output:
102	// a a b c d e f g h i j k l m n o p q r s t u v w x y z
103}
104