zclwrite: Body.FindAttribute implementation

zclwrite/ast.go

@@ -86,7 +86,15 @@ func (n *Body) AppendUnstructuredTokens(seq *TokenSeq) {
 // other calls may contain additional matching attributes that cannot be seen
 // by this method.
 func (n *Body) FindAttribute(name string) *Attribute {
-	panic("Body.FindAttribute not yet implemented")
+	nameBytes := []byte(name)
+	for _, item := range n.Items {
+		if attr, ok := item.(*Attribute); ok {
+			if attr.NameTokens.IsIdent(nameBytes) {
+				return attr
+			}
+		}
+	}
+	return nil
 }
 
 // SetAttributeValue either replaces the expression of an existing attribute

zclwrite/ast_test.go

@@ -0,0 +1,97 @@
+package zclwrite
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+
+	"github.com/davecgh/go-spew/spew"
+	"github.com/zclconf/go-zcl/zcl"
+	"github.com/zclconf/go-zcl/zcl/zclsyntax"
+)
+
+func TestBodyFindAttribute(t *testing.T) {
+	tests := []struct {
+		src  string
+		name string
+		want *TokenSeq
+	}{
+		{
+			"",
+			"a",
+			nil,
+		},
+		{
+			"a = 1\n",
+			"a",
+			&TokenSeq{
+				Tokens{
+					{
+						Type:  zclsyntax.TokenIdent,
+						Bytes: []byte{'a'},
+					},
+				},
+			},
+		},
+		{
+			"a = 1\nb = 1\nc = 1\n",
+			"a",
+			&TokenSeq{
+				Tokens{
+					{
+						Type:  zclsyntax.TokenIdent,
+						Bytes: []byte{'a'},
+					},
+				},
+			},
+		},
+		{
+			"a = 1\nb = 1\nc = 1\n",
+			"b",
+			&TokenSeq{
+				Tokens{
+					{
+						Type:  zclsyntax.TokenIdent,
+						Bytes: []byte{'b'},
+					},
+				},
+			},
+		},
+		{
+			"a = 1\nb = 1\nc = 1\n",
+			"c",
+			&TokenSeq{
+				Tokens{
+					{
+						Type:  zclsyntax.TokenIdent,
+						Bytes: []byte{'c'},
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%s in %s", test.name, test.src), func(t *testing.T) {
+			f, diags := ParseConfig([]byte(test.src), "", zcl.Pos{Line: 1, Column: 1})
+			if len(diags) != 0 {
+				for _, diag := range diags {
+					t.Logf("- %s", diag.Error())
+				}
+				t.Fatalf("unexpected diagnostics")
+			}
+
+			attr := f.Body.FindAttribute(test.name)
+			if attr == nil {
+				if test.want != nil {
+					t.Errorf("attribute found, but expecting not found")
+				}
+			} else {
+				got := attr.NameTokens
+				if !reflect.DeepEqual(got, test.want) {
+					t.Errorf("wrong result\ngot:  %s\nwant: %s", spew.Sdump(got), spew.Sdump(test.want))
+				}
+			}
+		})
+	}
+}

zclwrite/tokens.go

@@ -135,6 +135,35 @@ func (ts *TokenSeq) WriteTo(wr io.Writer) (int, error) {
 	return n, err
 }
 
+// SoloToken returns the single token represented by the receiving sequence,
+// or nil if the sequence does not represent exactly one token.
+func (ts *TokenSeq) SoloToken() *Token {
+	var ret *Token
+	found := false
+	ts.EachToken(func(tok *Token) {
+		if ret == nil && !found {
+			ret = tok
+			found = true
+		} else if ret != nil && found {
+			ret = nil
+		}
+	})
+	return ret
+}
+
+// IsIdent returns true if and only if the token sequence represents a single
+// ident token whose name matches the given string.
+func (ts *TokenSeq) IsIdent(name []byte) bool {
+	tok := ts.SoloToken()
+	if tok == nil {
+		return false
+	}
+	if tok.Type != zclsyntax.TokenIdent {
+		return false
+	}
+	return bytes.Equal(tok.Bytes, name)
+}
+
 // TokenSeqEmpty is a TokenSeq that contains no tokens. It can be used anywhere,
 // but its primary purpose is to be assigned as a replacement for a non-empty
 // TokenSeq when eliminating a section of an input file.