walk options of enumfield #140 (#141)

Co-authored-by: Ernest Micklei <ernest@ag5.com>

Author: emicklei

enum.go

@@ -140,6 +140,18 @@ type EnumField struct {
 	Parent        Visitee
 }
 
+// elements is part of elementContainer
+func (f *EnumField) elements() []Visitee {
+	return f.Elements
+}
+
+// takeLastComment is part of elementContainer
+// removes and returns the last element of the list if it is a Comment.
+func (f *EnumField) takeLastComment(expectedOnLine int) (last *Comment) {
+	last, f.Elements = takeLastCommentIfEndsOnLine(f.Elements, expectedOnLine)
+	return
+}
+
 // Accept dispatches the call to the visitor.
 func (f *EnumField) Accept(v Visitor) {
 	v.VisitEnumField(f)

enum_test.go

@@ -23,7 +23,9 @@
 
 package proto
 
-import "testing"
+import (
+	"testing"
+)
 
 func TestEnum(t *testing.T) {
 	proto := `
@@ -191,3 +193,32 @@ func TestEnumInlineCommentBeforeBody(t *testing.T) {
 		t.Errorf("got %d want %d lines", got, want)
 	}
 }
+
+func TestEnumFieldWalkWithComment(t *testing.T) {
+	src := `enum HideIt
+	  {
+		  PRIVATE = 1 [
+			// hidden
+			// field
+			(google.api.value_visibility).restriction = "HIDDEN"
+		  ];
+	  }
+	`
+	p := newParserOn(src)
+	e := new(Enum)
+	p.next()
+	if err := e.parse(p); err != nil {
+		t.Fatal(err)
+	}
+	var name, msg string
+	walk(e, WithOption(func(o *Option) {
+		name = o.Name
+		msg = o.Comment.Message()
+	}))
+	if got, want := name, "(google.api.value_visibility).restriction"; got != want {
+		t.Errorf("got [%v]:%T want [%v]:%T", got, got, want, want)
+	}
+	if got, want := msg, " hidden"; got != want {
+		t.Errorf("got [%v]:%T want [%v]:%T", got, got, want, want)
+	}
+}