From 1b37ff8aebb771f6f6c1c235f7446be33bde1652 Mon Sep 17 00:00:00 2001
From: Leslie Qi Wang <qiwa@pensando.io>
Date: Mon, 28 Nov 2022 14:50:53 -0800
Subject: [PATCH] add prefix option to allow caller for flattern customization

Signed-off-by: Leslie Qi Wang <qiwa@pensando.io>
---
 flat.go      |   6 ++-
 flat_test.go | 120 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+), 1 deletion(-)

diff --git a/flat.go b/flat.go
index 2a54476..c30dd2d 100644
--- a/flat.go
+++ b/flat.go
@@ -11,6 +11,7 @@ import (
 // Options the flatten options.
 // By default: Delimiter = "."
 type Options struct {
+	Prefix    string
 	Delimiter string
 	Safe      bool
 	MaxDepth  int
@@ -27,7 +28,7 @@ func Flatten(nested map[string]interface{}, opts *Options) (m map[string]interfa
 		}
 	}
 
-	m, err = flatten("", 0, nested, opts)
+	m, err = flatten(opts.Prefix, 0, nested, opts)
 
 	return
 }
@@ -125,6 +126,9 @@ func unflatten(flat map[string]interface{}, opts *Options) (nested map[string]in
 func uf(k string, v interface{}, opts *Options) (n interface{}) {
 	n = v
 
+	if opts.Prefix != "" {
+		k = strings.TrimPrefix(k, opts.Prefix+opts.Delimiter)
+	}
 	keys := strings.Split(k, opts.Delimiter)
 
 	for i := len(keys) - 1; i >= 0; i-- {
diff --git a/flat_test.go b/flat_test.go
index a73606a..c2b85e9 100644
--- a/flat_test.go
+++ b/flat_test.go
@@ -388,3 +388,123 @@ func TestUnflatten(t *testing.T) {
 		}
 	}
 }
+
+func TestFlattenPrefix(t *testing.T) {
+	tests := []struct {
+		given   string
+		options *Options
+		want    map[string]interface{}
+	}{
+		// test with different primitives
+		// String: 'world',
+		// Number: 1234.99,
+		// Boolean: true,
+		// null: null,
+		{
+			`{"hello": "world"}`,
+			&Options{Prefix: "test", Delimiter: "."},
+			map[string]interface{}{"test.hello": "world"},
+		},
+		{
+			`{"hello": 1234.99}`,
+			&Options{Prefix: "test", Delimiter: "_"},
+			map[string]interface{}{"test_hello": 1234.99},
+		},
+		{
+			`{"hello": true}`,
+			&Options{Prefix: "test", Delimiter: "-"},
+			map[string]interface{}{"test-hello": true},
+		},
+		{
+			`{"hello":{"world":"good morning"}}`,
+			&Options{Prefix: "test", Delimiter: "."},
+			map[string]interface{}{"test.hello.world": "good morning"},
+		},
+		{
+			`{"hello":{"world":1234.99}}`,
+			&Options{Prefix: "test", Delimiter: "_"},
+			map[string]interface{}{"test_hello_world": 1234.99},
+		},
+		{
+			`{"hello":{"world":true}}`,
+			&Options{Prefix: "test", Delimiter: "-"},
+			map[string]interface{}{"test-hello-world": true},
+		},
+	}
+	for i, test := range tests {
+		var given interface{}
+		err := json.Unmarshal([]byte(test.given), &given)
+		if err != nil {
+			t.Errorf("%d: failed to unmarshal test: %v", i+1, err)
+		}
+		got, err := Flatten(given.(map[string]interface{}), test.options)
+		if err != nil {
+			t.Errorf("%d: failed to flatten: %v", i+1, err)
+		}
+		if !reflect.DeepEqual(got, test.want) {
+			t.Errorf("%d: mismatch, got: %v want: %v", i+1, got, test.want)
+		}
+	}
+}
+
+func TestUnflattenPrefix(t *testing.T) {
+	tests := []struct {
+		flat    map[string]interface{}
+		options *Options
+		want    map[string]interface{}
+	}{
+		{
+			map[string]interface{}{"test.hello": "world"},
+			&Options{Prefix: "test", Delimiter: "."},
+			map[string]interface{}{"hello": "world"},
+		},
+		{
+			map[string]interface{}{"test_hello": 1234.56},
+			&Options{Prefix: "test", Delimiter: "_"},
+			map[string]interface{}{"hello": 1234.56},
+		},
+		{
+			map[string]interface{}{"test-hello": true},
+			&Options{Prefix: "test", Delimiter: "-"},
+			map[string]interface{}{"hello": true},
+		},
+		// nested twice
+		{
+			map[string]interface{}{"test.hello.world.again": "good morning"},
+			&Options{Prefix: "test", Delimiter: "."},
+			map[string]interface{}{
+				"hello": map[string]interface{}{
+					"world": map[string]interface{}{
+						"again": "good morning",
+					},
+				},
+			},
+		},
+		// custom delimiter
+		{
+			map[string]interface{}{
+				"test hello world again": "good morning",
+			},
+			&Options{
+				Prefix:    "test",
+				Delimiter: " ",
+			},
+			map[string]interface{}{
+				"hello": map[string]interface{}{
+					"world": map[string]interface{}{
+						"again": "good morning",
+					},
+				},
+			},
+		},
+	}
+	for i, test := range tests {
+		got, err := Unflatten(test.flat, test.options)
+		if err != nil {
+			t.Errorf("%d: failed to unflatten: %v", i+1, err)
+		}
+		if !reflect.DeepEqual(got, test.want) {
+			t.Errorf("%d: mismatch, got: %v want: %v", i+1, got, test.want)
+		}
+	}
+}