github.com/k14s/starlark-go@v0.0.0-20200720175618-3a5c849cc368/syntax/walk_test.go (about)

     1  package syntax_test
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"log"
     7  	"reflect"
     8  	"strings"
     9  	"testing"
    10  
    11  	"github.com/k14s/starlark-go/syntax"
    12  )
    13  
    14  func TestWalk(t *testing.T) {
    15  	const src = `
    16  for 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 := `
    46  File
    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.
    73  func ExampleWalk() {
    74  	const src = `
    75  load("library", "a")
    76  
    77  def b(c, *, d=e):
    78      f += {g: h}
    79      i = -(j)
    80      return k.l[m + n]
    81  
    82  for 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  }