From aa708eb4decbd5699aca09a006d0b4a56706a2d0 Mon Sep 17 00:00:00 2001
From: Brandon Buck <lordizuriel@gmail.com>
Date: Thu, 30 Apr 2015 22:39:43 -0500
Subject: [PATCH] Clean up, remove zero as 'empty' and add 'omitzero' option

---
 encode.go      | 37 +++++++++++++++++++++++++------------
 encode_test.go | 22 ++++++++++++++++++++++
 2 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/encode.go b/encode.go
index 63be32e2..c7e227c7 100644
--- a/encode.go
+++ b/encode.go
@@ -344,17 +344,10 @@ func (enc *Encoder) eStruct(key Key, rv reflect.Value) {
 				keyName = sft.Name
 			}
 
-			opts := strings.Split(keyName, ",")
-			keyName = opts[0]
-			omitEmpty := false
-			if len(opts) > 1 {
-				for _, opt := range opts[1:] {
-					if opt == "omitempty" {
-						omitEmpty = true
-					}
-				}
-			}
-			if omitEmpty && isEmpty(sf) {
+			keyName, opts := getOptions(keyName)
+			if _, ok := opts["omitempty"]; ok && isEmpty(sf) {
+				continue
+			} else if _, ok := opts["omitzero"]; ok && isZero(sf) {
 				continue
 			}
 
@@ -448,7 +441,20 @@ func tomlArrayType(rv reflect.Value) tomlType {
 	return firstType
 }
 
-func isEmpty(rv reflect.Value) bool {
+func getOptions(keyName string) (string, map[string]struct{}) {
+	opts := make(map[string]struct{})
+	ss := strings.Split(keyName, ",")
+	name := ss[0]
+	if len(ss) > 1 {
+		for _, opt := range ss {
+			opts[opt] = struct{}{}
+		}
+	}
+
+	return name, opts
+}
+
+func isZero(rv reflect.Value) bool {
 	switch rv.Kind() {
 	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
 		if rv.Int() == 0 {
@@ -462,6 +468,13 @@ func isEmpty(rv reflect.Value) bool {
 		if rv.Float() == 0.0 {
 			return true
 		}
+	}
+
+	return false
+}
+
+func isEmpty(rv reflect.Value) bool {
+	switch rv.Kind() {
 	case reflect.String:
 		if len(strings.TrimSpace(rv.String())) == 0 {
 			return true
diff --git a/encode_test.go b/encode_test.go
index 2f30d5cf..445ca8e2 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -469,6 +469,28 @@ func TestEncodeWithOmitEmpty(t *testing.T) {
 	encodeExpected(t, "simple with omitempty, not empty", value, expected, nil)
 }
 
+func TestEncodeWithOmitZero(t *testing.T) {
+	type simple struct {
+		Number   int     `toml:"number,omitzero"`
+		Real     float64 `toml:"real,omitzero"`
+		Unsigned uint    `toml:"unsigned,omitzero"`
+	}
+
+	value := simple{0, 0.0, uint(0)}
+	expected := ""
+
+	encodeExpected(t, "simple with omitzero, all zero", value, expected, nil)
+
+	value.Number = 10
+	value.Real = 20
+	value.Unsigned = 5
+	expected = `number = 10
+real = 20.0
+unsigned = 5
+`
+	encodeExpected(t, "simple with omitzero, non-zero", value, expected, nil)
+}
+
 func encodeExpected(
 	t *testing.T, label string, val interface{}, wantStr string, wantErr error,
 ) {