From 3abdf581913ef08fbc69ba016e4cb61628faee17 Mon Sep 17 00:00:00 2001
From: Jason Lee <huacnlee@gmail.com>
Date: Tue, 1 Jun 2021 00:07:47 +0800
Subject: [PATCH] Format, FormatHTML, Unformat, UnformatHTML support custom
 option for addition processes

---
 format.go        | 10 +++++++++-
 format_test.go   | 20 ++++++++++++++++++++
 html.go          | 12 ++++++++----
 unformat.go      | 12 +++++++++++-
 unformat_test.go | 19 +++++++++++++++++++
 5 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/format.go b/format.go
index 142b1d7..38bd852 100644
--- a/format.go
+++ b/format.go
@@ -61,8 +61,12 @@ func spaceDashWithHans(in string) (out string) {
 	return out
 }
 
+type Option interface {
+	Format(text string) string
+}
+
 // Format auto format string to add spaces between Chinese and English words.
-func Format(in string) (out string) {
+func Format(in string, options ...Option) (out string) {
 	out = in
 
 	out = halfwidth(out)
@@ -75,5 +79,9 @@ func Format(in string) (out string) {
 	out = removeFullDateSpacing(out)
 	out = spaceDashWithHans(out)
 
+	for _, opt := range options {
+		out = opt.Format(out)
+	}
+
 	return
 }
diff --git a/format_test.go b/format_test.go
index d7bb4f6..c48b2c8 100644
--- a/format_test.go
+++ b/format_test.go
@@ -3,6 +3,8 @@ package autocorrect
 import (
 	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func assertEqual(t *testing.T, exptected, actual string) {
@@ -61,6 +63,24 @@ func TestFormat(t *testing.T) {
 	assertCases(t, cases)
 }
 
+type customFormat struct{}
+
+func (c customFormat) Format(text string) string {
+	return strings.ReplaceAll(text, "AAAA", "BBBB")
+}
+
+func TestFormatWithOptions(t *testing.T) {
+	assert.Equal(t, "测试 options 你好 BBBB", Format("测试options你好AAAA", customFormat{}))
+}
+
+func Test_FormatHTMLWithOptions(t *testing.T) {
+	out, err := FormatHTML("<p>测试options你好AAAA</p>", customFormat{})
+	if err != nil {
+		t.Error(err)
+	}
+	assert.Equal(t, "<p>测试 options 你好 BBBB</p>", out)
+}
+
 func TestFormatForSpecialChars(t *testing.T) {
 	cases := map[string]string{
 		"记事本,记事本显示阅读次数#149":           "记事本，记事本显示阅读次数#149",
diff --git a/html.go b/html.go
index 6ebbe82..7183d43 100644
--- a/html.go
+++ b/html.go
@@ -16,13 +16,17 @@ var (
 )
 
 // FormatHTML format HTML content
-func FormatHTML(body string) (out string, err error) {
-	return processHTML(body, Format)
+func FormatHTML(body string, options ...Option) (out string, err error) {
+	return processHTML(body, func(text string) string {
+		return Format(text, options...)
+	})
 }
 
 // UnformatHTML cleanup spaces for HTML
-func UnformatHTML(body string) (out string, err error) {
-	return processHTML(body, Unformat)
+func UnformatHTML(body string, options ...UnformatOption) (out string, err error) {
+	return processHTML(body, func(text string) string {
+		return Unformat(text, options...)
+	})
 }
 
 func processHTML(body string, fn func(plainText string) string) (out string, err error) {
diff --git a/unformat.go b/unformat.go
index d2025d7..ba04940 100644
--- a/unformat.go
+++ b/unformat.go
@@ -6,8 +6,18 @@ var (
 	removeSpaceRe = regexp.MustCompile(`(` + spaceRe + `+)?(` + cjk + `)(` + spaceRe + `+)?`)
 )
 
+// UnformatOption addition unformat options
+type UnformatOption interface {
+	Unformat(text string) string
+}
+
 // Unformat to remove all spaces
-func Unformat(text string) string {
+func Unformat(text string, options ...UnformatOption) string {
 	text = removeSpaceRe.ReplaceAllString(text, "$2")
+
+	for _, opt := range options {
+		text = opt.Unformat(text)
+	}
+
 	return text
 }
diff --git a/unformat_test.go b/unformat_test.go
index 9e2615d..035b6b9 100644
--- a/unformat_test.go
+++ b/unformat_test.go
@@ -1,6 +1,7 @@
 package autocorrect
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -15,3 +16,21 @@ func TestUnformat(t *testing.T) {
 	expected := "据港交所最新权益披露资料显示，2019年12月27日，三生制药获JP Morgan Chase & Co.每股均价9.582港元，增持270.3万股，总价约2590万港元。"
 	assert.Equal(t, expected, Unformat(raw))
 }
+
+type customUnformat struct{}
+
+func (c customUnformat) Unformat(text string) string {
+	return strings.ReplaceAll(text, "BBBB", "AAAA")
+}
+
+func TestUnformatWithOptions(t *testing.T) {
+	assert.Equal(t, "增持270.3万股AAAA", Unformat("增持270.3万股BBBB", customUnformat{}))
+}
+
+func Test_UnformatHTMLWithOptions(t *testing.T) {
+	out, err := UnformatHTML("<p>测试 options 你好 BBBB</p>", customUnformat{})
+	if err != nil {
+		t.Error(err)
+	}
+	assert.Equal(t, "<p>测试options你好AAAA</p>", out)
+}