extract messages from slices and maps (#169)

2 files changed, 106 insertions, 14 deletions

diff --git a/v2/goi18n/extract_command.go b/v2/goi18n/extract_command.go
index 0ea8411..cd01e72 100644
--- a/v2/goi18n/extract_command.go
+++ b/v2/goi18n/extract_command.go
@@ -140,29 +140,78 @@ func (e *extractor) err(err error) {
 }
 
 func (e *extractor) Visit(node ast.Node) ast.Visitor {
-	e.extractMessage(node)
+	e.extractMessages(node)
 	return e
 }
 
-func (e *extractor) extractMessage(node ast.Node) {
+func (e *extractor) extractMessages(node ast.Node) {
 	cl, ok := node.(*ast.CompositeLit)
 	if !ok {
 		return
 	}
-	se, ok := cl.Type.(*ast.SelectorExpr)
-	if !ok {
-		return
+	switch t := cl.Type.(type) {
+	case *ast.SelectorExpr:
+		if !e.isMessageType(t) {
+			return
+		}
+		e.extractMessage(cl)
+	case *ast.ArrayType:
+		if !e.isMessageType(t.Elt) {
+			return
+		}
+		for _, el := range cl.Elts {
+			ecl, ok := el.(*ast.CompositeLit)
+			if !ok {
+				continue
+			}
+			e.extractMessage(ecl)
+		}
+	case *ast.MapType:
+		if !e.isMessageType(t.Value) {
+			return
+		}
+		for _, el := range cl.Elts {
+			kve, ok := el.(*ast.KeyValueExpr)
+			if !ok {
+				continue
+			}
+			vcl, ok := kve.Value.(*ast.CompositeLit)
+			if !ok {
+				continue
+			}
+			e.extractMessage(vcl)
+		}
+	}
+}
+
+func (e *extractor) isMessageType(expr ast.Expr) bool {
+	se := unwrapSelectorExpr(expr)
+	if se == nil {
+		return false
 	}
 	if se.Sel.Name != "Message" && se.Sel.Name != "LocalizeConfig" {
-		return
+		return false
 	}
 	x, ok := se.X.(*ast.Ident)
 	if !ok {
-		return
+		return false
 	}
-	if x.Name != e.i18nPackageName {
-		return
+	return x.Name == e.i18nPackageName
+}
+
+func unwrapSelectorExpr(e ast.Expr) *ast.SelectorExpr {
+	switch et := e.(type) {
+	case *ast.SelectorExpr:
+		return et
+	case *ast.StarExpr:
+		se, _ := et.X.(*ast.SelectorExpr)
+		return se
+	default:
+		return nil
 	}
+}
+
+func (e *extractor) extractMessage(cl *ast.CompositeLit) {
 	data := make(map[string]string)
 	for _, elt := range cl.Elts {
 		kve, ok := elt.(*ast.KeyValueExpr)
@@ -182,11 +231,10 @@ func (e *extractor) extractMessage(node ast.Node) {
 	if len(data) == 0 {
 		return
 	}
-	if se.Sel.Name == "Message" {
-		e.messages = append(e.messages, internal.MustNewMessage(data))
-	} else if messageID := data["MessageID"]; messageID != "" {
-		e.messages = append(e.messages, &i18n.Message{ID: messageID})
+	if messageID := data["MessageID"]; messageID != "" {
+		data["ID"] = messageID
 	}
+	e.messages = append(e.messages, internal.MustNewMessage(data))
 }
 
 func extractStringLiteral(expr ast.Expr) (string, bool) {
diff --git a/v2/goi18n/extract_command_test.go b/v2/goi18n/extract_command_test.go
index 2e3ae94..4aacafa 100644
--- a/v2/goi18n/extract_command_test.go
+++ b/v2/goi18n/extract_command_test.go
@@ -54,6 +54,50 @@ b = "a \" b"
 `),
 		},
 		{
+			name:     "array",
+			fileName: "file.go",
+			file: `package main
+
+			import "github.com/nicksnyder/go-i18n/v2/i18n"
+
+			var a = []*i18n.Message{
+				{
+					ID:    "a",
+					Other: "a",
+				},
+				{
+					ID:    "b",
+					Other: "b",
+				},
+			}
+			`,
+			activeFile: []byte(`a = "a"
+b = "b"
+`),
+		},
+		{
+			name:     "map",
+			fileName: "file.go",
+			file: `package main
+
+			import "github.com/nicksnyder/go-i18n/v2/i18n"
+
+			var a = map[string]*i18n.Message{
+				"a": {
+					ID:    "a",
+					Other: "a",
+				},
+				"b": {
+					ID:    "b",
+					Other: "b",
+				},
+			}
+			`,
+			activeFile: []byte(`a = "a"
+b = "b"
+`),
+		},
+		{
 			name:     "no extract from test",
 			fileName: "file_test.go",
 			file: `package main
@@ -144,7 +188,7 @@ zero = "Zero translation"
 	}
 
 	for _, test := range tests {
-		t.Run(test.name+" active file", func(t *testing.T) {
+		t.Run(test.name, func(t *testing.T) {
 			indir := mustTempDir("TestExtractCommandIn")
 			defer os.RemoveAll(indir)
 			outdir := mustTempDir("TestExtractCommandOut")