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