2 files changed, 37 insertions, 22 deletions

diff --git a/cmd/gettext/gettext.go b/cmd/gettext/gettext.go
index 2dbd65e..dd1ba0c 100644
--- a/cmd/gettext/gettext.go
+++ b/cmd/gettext/gettext.go
@@ -3,6 +3,9 @@ package main
 import (
 	"flag"
 	"fmt"
+	"go/ast"
+	"go/parser"
+	"go/token"
 	"io/fs"
 	"os"
 	"path/filepath"
@@ -80,14 +83,25 @@ func (p ArgsParser) ArgsCount() int {
 
 func extractKeys(content string) map[string]int {
 	extractedKeys := make(map[string]int, 0)
-	reg := regexp.MustCompile(`\{\{\s*call \$?\.T\s*"(\w+)"([\w.\s",:()]*)}}`)
+	reg := regexp.MustCompile(`\{ ?(login\.Tr\("\w+"[\w"(),. :]*\) ?)\}`)
 	matches := reg.FindAllStringSubmatch(content, -1)
 	for _, match := range matches {
-		key := match[1]
-		args := strings.TrimSpace(match[2])
-		parser := ArgsParser{}
-		parser.Parse(args)
-		extractedKeys[key] = parser.ArgsCount()
+		fakePackage := fmt.Sprintf(`package main
+	
+		func main() {
+			%s
+		}
+		`, match[1])
+		fset := token.NewFileSet()
+		f, err := parser.ParseFile(fset, "", fakePackage, 0)
+		kcore.Expect(err, "error parsing file")
+		main := f.Decls[0].(*ast.FuncDecl)
+		call := main.Body.List[0].(*ast.ExprStmt).X.(*ast.CallExpr)
+		selector := call.Fun.(*ast.SelectorExpr)
+		if selector.Sel.Name == "Tr" {
+			key := strings.Trim(call.Args[0].(*ast.BasicLit).Value, `"`)
+			extractedKeys[key] = len(call.Args) - 1
+		}
 	}
 	return extractedKeys
 }
@@ -95,8 +109,8 @@ func extractKeys(content string) map[string]int {
 func extractAllKeys() map[string]int {
 	extractedKeys := make(map[string]int, 0)
 
-	err := filepath.Walk("templates", func(path string, info fs.FileInfo, err error) error {
-		if !info.IsDir() && filepath.Ext(path) == ".html" {
+	err := filepath.Walk(".", func(path string, info fs.FileInfo, err error) error {
+		if !info.IsDir() && filepath.Ext(path) == ".templ" {
 			content, err := os.ReadFile(path) // #nosec G304
 			kcore.Expect(err, "error reading file")
 			for key, value := range extractKeys(string(content)) {
diff --git a/cmd/gettext/gettext_test.go b/cmd/gettext/gettext_test.go
index f1947be..8952c73 100644
--- a/cmd/gettext/gettext_test.go
+++ b/cmd/gettext/gettext_test.go
@@ -7,7 +7,7 @@ import (
 )
 
 func TestExtractKeys(t *testing.T) {
-	content := `{{ call .T "Hello" }}`
+	content := `{ login.Tr("Hello") }`
 
 	keys := extractKeys(content)
 
@@ -16,16 +16,7 @@ func TestExtractKeys(t *testing.T) {
 }
 
 func TestExtractKeysWithoutSpaces(t *testing.T) {
-	content := `{{call .T "Hello"}}`
-
-	keys := extractKeys(content)
-
-	assert.Len(t, keys, 1)
-	assert.Equal(t, 0, keys["Hello"])
-}
-
-func TestExtractKeysWithDollar(t *testing.T) {
-	content := `{{ call $.T "Hello" }}`
+	content := `{login.Tr("Hello")}`
 
 	keys := extractKeys(content)
 
@@ -34,7 +25,7 @@ func TestExtractKeysWithDollar(t *testing.T) {
 }
 
 func TestExtractKeysWithOneArgs(t *testing.T) {
-	content := `{{ call .T "approveButton" .Username }}`
+	content := `{ login.Tr("approveButton", a.Username) }`
 
 	keys := extractKeys(content)
 
@@ -43,7 +34,7 @@ func TestExtractKeysWithOneArgs(t *testing.T) {
 }
 
 func TestExtractKeysWithSeveralArgs(t *testing.T) {
-	content := `{{ call .T "approveButton" .Username .Age .Email }}`
+	content := `{ login.Tr("approveButton", a.Username, a.Age, a.Email) }`
 
 	keys := extractKeys(content)
 
@@ -52,10 +43,20 @@ func TestExtractKeysWithSeveralArgs(t *testing.T) {
 }
 
 func TestExtractKeysWithComplexArgs(t *testing.T) {
-	content := `{{ call $.T "raceStart_chosen" (.StartAt.Format "Monday, January 2, 2006 at 15:04") }}`
+	content := `{ login.Tr("raceStart_chosen", a.StartAt.Format("Monday, January 2, 2006 at 15:04")) }`
 
 	keys := extractKeys(content)
 
 	assert.Len(t, keys, 1)
 	assert.Equal(t, 1, keys["raceStart_chosen"])
 }
+
+func TestExtractMultipleKeys(t *testing.T) {
+	content := `{ login.Tr("test_1", a.StartAt.Format("Monday, January 2, 2006 at 15:04")) }{ login.Tr("test_2", a.Test) }`
+
+	keys := extractKeys(content)
+
+	assert.Len(t, keys, 2)
+	assert.Equal(t, 1, keys["test_1"])
+	assert.Equal(t, 1, keys["test_2"])
+}