Refine colorhttp func

gavv · Oct 3, 2023 · e6879c0 · e6879c0
1 parent 420f3ae
commit e6879c0
Showing 1 changed file with 74 additions and 64 deletions.
diff --git a/formatter.go b/formatter.go
@@ -1009,63 +1009,104 @@ var defaultTemplateFuncs = template.FuncMap{
 		}
 		return color.New(colorAttr).Sprint(input)
 	},
-	"colorhttp": func(enable bool, colorName string, isResponse bool, input string) string {
+	"colorhttp": func(enable bool, isResponse bool, input string) string {
 		if !enable {
 			return input
 		}
 
-		colorAttr := color.Reset
-		present := true
-		var sb strings.Builder
-		if ca, ok := defaultColors[colorName]; ok {
-			colorAttr = ca
-		}
+		methodColor := color.New(defaultColors["HiMagenta"])
+		statusColor := color.New(defaultColors["HiMagenta"])
+		headerColor := color.New(defaultColors["Cyan"])
 
-		httpMethod := map[string]bool{
-			"GET":    present,
-			"POST":   present,
-			"PATCH":  present,
-			"DELETE": present,
-			"PUT":    present,
-		}
-
-		getColor := func(word string, forceApplyColor bool) color.Attribute {
-			if v := httpMethod[word]; v == present {
-				return colorAttr
-			}
-			if word[len(word)-1] == ':' {
-				return colorAttr
-			}
-			if forceApplyColor {
-				return colorAttr
-			}
+		var sb strings.Builder
 
-			return color.Reset
-		}
 		isFirstLine := true
 		for _, line := range strings.Split(input, "\n") {
 			if sb.Len() != 0 {
 				sb.WriteString("\n")
 			}
 
-			words := strings.Fields(line)
+			line = strings.TrimSuffix(line, "\n")
+			line = strings.TrimSuffix(line, "\r")
+
+			var words []string
+			if isFirstLine {
+				words = strings.SplitN(line, " ", -1)
+			} else {
+				words = strings.SplitN(line, " ", 2)
+			}
+
 			wordLen := len(words)
 			for index, word := range words {
-				forceApplyColor := index != wordLen-1 && isResponse && isFirstLine
-				sb.WriteString(color.New(getColor(word, forceApplyColor)).Sprint(word))
+				var applyColor *color.Color
+
+				if isFirstLine {
+					if isResponse {
+						if index != 0 && index != wordLen-1 {
+							applyColor = statusColor
+						}
+					} else {
+						if index == 0 {
+							applyColor = methodColor
+						}
+					}
+				} else {
+					if index == 0 {
+						applyColor = headerColor
+					}
+				}
+
+				if word != "" && applyColor != nil {
+					sb.WriteString(applyColor.Sprint(word))
+				} else {
+					sb.WriteString(word)
+				}
+
 				sb.WriteString(" ")
 			}
+
 			isFirstLine = false
 		}
 
 		return sb.String()
 	},
+	"colorjson": func(enable bool, colorName, input string) string {
+		if !enable {
+			return input
+		}
+
+		fallbackColor := color.Reset
+		if attr, ok := defaultColors[colorName]; ok {
+			fallbackColor = attr
+		}
+
+		var parsedInput interface{}
+		err := json.Unmarshal([]byte(input), &parsedInput)
+		if err != nil {
+			return color.New(fallbackColor).Sprint(input)
+		}
+
+		formatter := colorjson.NewFormatter()
+		formatter.KeyColor = color.New(color.Reset)
+		formatter.StringColor = color.New(defaultColors["HiMagenta"])
+		formatter.NumberColor = color.New(defaultColors["Cyan"])
+		formatter.BoolColor = color.New(defaultColors["Cyan"])
+		formatter.NullColor = color.New(defaultColors["Cyan"])
+		formatter.Indent = 2
+
+		b, err := formatter.Marshal(parsedInput)
+		if err != nil {
+			return color.New(fallbackColor).Sprint(input)
+		}
+
+		return string(b)
+	},
 	"colordiff": func(enable bool, input string) string {
 		if !enable {
 			return input
 		}
 
-		prefixColor := []struct {
+		prefixColor := [...]struct {
 			prefix string
 			color  color.Attribute
 		}{
@@ -1096,37 +1137,6 @@ var defaultTemplateFuncs = template.FuncMap{
 
 		return sb.String()
 	},
-	"colorjson": func(enable bool, colorName, input string) string {
-		if !enable {
-			return input
-		}
-
-		fallbackColor := color.Reset
-		if attr, ok := defaultColors[colorName]; ok {
-			fallbackColor = attr
-		}
-
-		var parsedInput interface{}
-		err := json.Unmarshal([]byte(input), &parsedInput)
-		if err != nil {
-			return color.New(fallbackColor).Sprint(input)
-		}
-
-		formatter := colorjson.NewFormatter()
-		formatter.KeyColor = color.New(color.Reset)
-		formatter.StringColor = color.New(defaultColors["HiMagenta"])
-		formatter.NumberColor = color.New(defaultColors["Cyan"])
-		formatter.BoolColor = color.New(defaultColors["Cyan"])
-		formatter.NullColor = color.New(defaultColors["Cyan"])
-		formatter.Indent = 2
-
-		b, err := formatter.Marshal(parsedInput)
-		if err != nil {
-			return color.New(fallbackColor).Sprint(input)
-		}
-
-		return string(b)
-	},
 }
 
 var defaultSuccessTemplate = `[OK] {{ join .LineWidth .AssertPath }}`
@@ -1149,11 +1159,11 @@ request name: {{ .RequestName | color $.EnableColors "Cyan" }}
 {{- end -}}
 {{- if .HaveRequest }}
 
-request: {{ .Request | indent | trim | colorhttp $.EnableColors "HiMagenta" false}}
+request: {{ .Request | colorhttp $.EnableColors false | indent | trim }}
 {{- end -}}
 {{- if .HaveResponse }}
 
-response: {{ .Response | indent | trim | colorhttp $.EnableColors "HiMagenta" true}}
+response: {{ .Response | colorhttp $.EnableColors true | indent | trim }}
 {{- end -}}
 {{- if .HaveStacktrace }}