From 86fb472d0f1dec91aff30a09170c524acf2a068a Mon Sep 17 00:00:00 2001
From: Xuanwo <github@xuanwo.io>
Date: Fri, 3 Sep 2021 10:20:48 +0800
Subject: [PATCH] feat: Add omit wrap support

Signed-off-by: Xuanwo <github@xuanwo.io>
---
 README.md     | 32 ++++++++++++++++++++++++++++++++
 function.go   |  7 ++++++-
 group.go      | 26 ++++++++++++++++++++++++--
 group_test.go | 19 +++++++++++++++++++
 import.go     | 10 +++++++++-
 5 files changed, 90 insertions(+), 4 deletions(-)
 create mode 100644 group_test.go

diff --git a/README.md b/README.md
index 36ab1ba..b4e4323 100644
--- a/README.md
+++ b/README.md
@@ -2,6 +2,38 @@
 
 `gg` is a General Golang Code Generator: A Good Game to play with Golang.
 
+```go
+package main
+
+import (
+	"fmt"
+	
+	"github.com/Xuanwo/gg"
+)
+
+func main() {
+	f := Group()
+	f.Package("main")
+	f.Imports().Path("fmt")
+	f.Function("main").Body(
+		StringF(`fmt.Println("%s")`, "Hello, World!"),
+	)
+	fmt.Println(f.String())
+}
+```
+
+Output (after `go fmt`)
+
+```go
+package main
+
+import "fmt"
+
+func main() {
+	fmt.Println("Hello, World!")
+}
+```
+
 ## Acknowledgement
 
 - `gg` is inspired by [dave/jennifer](https://github.com/dave/jennifer), I borrowed most ideas and some code from it. Nice work!
diff --git a/function.go b/function.go
index 7033683..1fd5310 100644
--- a/function.go
+++ b/function.go
@@ -12,13 +12,18 @@ type ifunction struct {
 }
 
 func Function(name string) *ifunction {
-	return &ifunction{
+	i := &ifunction{
 		name:       name,
 		comments:   newGroup("", "", "\n"),
 		parameters: newGroup("(", ")", ","),
 		results:    newGroup("(", ")", ","),
 		body:       newGroup("{\n", "}", "\n"),
 	}
+	// We should omit the `()` if result is empty
+	i.results.omitWrapIf = func() bool {
+		return i.results.length() == 0
+	}
+	return i
 }
 
 func (f *ifunction) render(w io.Writer) {
diff --git a/group.go b/group.go
index aa485e4..0bddfd5 100644
--- a/group.go
+++ b/group.go
@@ -23,6 +23,20 @@ type group struct {
 	open      string
 	close     string
 	separator string
+
+	// If this result is true, we will omit the wrap like `()`, `{}`.
+	omitWrapIf func() bool
+}
+
+func (g *group) length() int {
+	return len(g.items)
+}
+
+func (g *group) shouldOmitWrap() bool {
+	if g.omitWrapIf == nil {
+		return false
+	}
+	return g.omitWrapIf()
 }
 
 func (g *group) append(node ...Node) *group {
@@ -31,7 +45,7 @@ func (g *group) append(node ...Node) *group {
 }
 
 func (g *group) render(w io.Writer) {
-	if g.open != "" {
+	if g.open != "" && !g.shouldOmitWrap() {
 		writeString(w, g.open)
 	}
 
@@ -44,7 +58,7 @@ func (g *group) render(w io.Writer) {
 		isfirst = false
 	}
 
-	if g.close != "" {
+	if g.close != "" && !g.shouldOmitWrap() {
 		writeString(w, g.close)
 	}
 }
@@ -62,6 +76,14 @@ func (g *group) WriteFile(path string) error {
 	return nil
 }
 
+func (g *group) String() string {
+	buf := pool.Get()
+	defer buf.Free()
+
+	g.render(buf)
+	return buf.String()
+}
+
 func (g *group) Comment(content string) *group {
 	g.append(Comment(content))
 	return g
diff --git a/group_test.go b/group_test.go
new file mode 100644
index 0000000..ab9f09c
--- /dev/null
+++ b/group_test.go
@@ -0,0 +1,19 @@
+package gg
+
+import "fmt"
+
+func ExampleGroup() {
+	f := Group()
+	f.Package("main")
+	f.Imports().Path("fmt")
+	f.Function("main").Body(
+		StringF(`fmt.Println("%s")`, "Hello, World!"),
+	)
+	fmt.Println(f.String())
+	// Output:
+	// package main
+	//
+	// import "fmt"
+	// func main(){
+	// fmt.Println("Hello, World!")}
+}
diff --git a/import.go b/import.go
index 7e890d2..c19f042 100644
--- a/import.go
+++ b/import.go
@@ -7,12 +7,20 @@ type iimport struct {
 }
 
 func Imports() *iimport {
-	return &iimport{
+	i := &iimport{
 		items: newGroup("(", ")", "\n"),
 	}
+	i.items.omitWrapIf = func() bool {
+		return i.items.length() <= 1
+	}
+	return i
 }
 
 func (i *iimport) render(w io.Writer) {
+	// Don't need to render anything if import is empty
+	if i.items.length() == 0 {
+		return
+	}
 	writeString(w, "import ")
 	i.items.render(w)
 }