symtable: add Find() method

ncw · ncw · commit 329523c768da · 2015-04-30T17:10:37.000+01:00
diff --git a/symtable/symtable.go b/symtable/symtable.go
@@ -12,6 +12,7 @@ package symtable
 import (
 	"fmt"
 	"log"
+	"sort"
 	"strings"
 
 	"github.com/ncw/gpython/ast"
@@ -782,3 +783,15 @@ func (st *SymTable) Analyze() {
 	global := make(StringSet)
 	st.AnalyzeBlock(nil, free, global)
 }
+
+// Return a sorted list of symbol names if the scope of a name matches
+// either scopeType or flag is set
+func (st *SymTable) Find(scopeType Scope, flag DefUseFlags) (out []string) {
+	for name, v := range st.Symbols {
+		if v.Scope == scopeType || (v.Flags&flag) != 0 {
+			out = append(out, name)
+		}
+	}
+	sort.Strings(out)
+	return out
+}
diff --git a/symtable/symtable_test.go b/symtable/symtable_test.go
@@ -169,3 +169,44 @@ func TestStringer(t *testing.T) {
 	EqString(t, "BlockType", "ClassBlock", ClassBlock.String())
 	EqString(t, "BlockType", "BlockType(100)", BlockType(100).String())
 }
+
+func TestSymTableFind(t *testing.T) {
+	st := &SymTable{
+		Symbols: Symbols{
+			"x": Symbol{
+				Flags: DefLocal | DefNonlocal,
+				Scope: ScopeFree,
+			},
+			"a": Symbol{
+				Flags: DefLocal | DefNonlocal,
+				Scope: ScopeFree,
+			},
+			"b": Symbol{
+				Flags: DefLocal | DefNonlocal,
+				Scope: ScopeFree,
+			},
+			"c": Symbol{
+				Flags: DefNonlocal,
+				Scope: ScopeCell,
+			},
+			"d": Symbol{
+				Flags: DefLocal,
+				Scope: ScopeCell,
+			},
+		},
+	}
+
+	for _, test := range []struct {
+		scope Scope
+		flag  DefUseFlags
+		want  []string
+	}{
+		{scope: ScopeGlobalExplicit, flag: 0, want: []string{}},
+		{scope: ScopeFree, flag: 0, want: []string{"a", "b", "x"}},
+		{scope: ScopeFree, flag: DefLocal, want: []string{"a", "b", "d", "x"}},
+		{scope: 0, flag: DefNonlocal, want: []string{"a", "b", "c", "x"}},
+	} {
+		got := st.Find(test.scope, test.flag)
+		EqStrings(t, fmt.Sprintf("Scope %v, Flag %v", test.scope, test.flag), test.want, got)
+	}
+}
