From 648297512ebd44f8ad7b65cad39a9714969744d6 Mon Sep 17 00:00:00 2001
From: shixzie <shixzie@gmail.com>
Date: Tue, 27 Jun 2017 11:03:53 -0400
Subject: [PATCH] time.Duration support

---
 README.md   |  1 +
 nlp.go      |  6 ++++++
 nlp_test.go | 22 ++++++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/README.md b/README.md
index bc623f4..87b307d 100644
--- a/README.md
+++ b/README.md
@@ -15,6 +15,7 @@ uint uint8 uint16 uint32 uint64
 float32 float64
 string
 time.Time
+time.Duration
 ```
 
 ## Installation
diff --git a/nlp.go b/nlp.go
index 9ea3ba4..a954294 100644
--- a/nlp.go
+++ b/nlp.go
@@ -141,6 +141,9 @@ func (nl *NL) RegisterModel(i interface{}, samples []string, ops ...ModelOption)
 			if v, ok := val.Field(i).Interface().(time.Time); ok {
 				mod.fields = append(mod.fields, field{i, tpy.Field(i).Name, v})
 				continue NextField
+			} else if v, ok := val.Field(i).Interface().(time.Duration); ok {
+				mod.fields = append(mod.fields, field{i, tpy.Field(i).Name, v})
+				continue NextField
 			}
 			switch val.Field(i).Kind() {
 			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Float32, reflect.Float64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.String:
@@ -305,6 +308,9 @@ func (m *model) fit(expr string) interface{} {
 			case time.Time:
 				v, _ := time.ParseInLocation(m.timeFormat, e.value, m.timeLocation)
 				val.Elem().Field(e.field.index).Set(reflect.ValueOf(v))
+			case time.Duration:
+				v, _ := time.ParseDuration(e.value)
+				val.Elem().Field(e.field.index).Set(reflect.ValueOf(v))
 			}
 		}
 	}
diff --git a/nlp_test.go b/nlp_test.go
index 9e52944..58401fe 100644
--- a/nlp_test.go
+++ b/nlp_test.go
@@ -16,6 +16,7 @@ func TestNL_P(t *testing.T) {
 		Uint   uint
 		Float  float32
 		Time   time.Time
+		Dur    time.Duration
 	}
 
 	tSamples := []string{
@@ -24,6 +25,7 @@ func TestNL_P(t *testing.T) {
 		"uint {Uint}",
 		"float {Float}",
 		"time {Time}",
+		"dur {Dur}",
 		"string {String} int {Int}",
 		"string {String} time {Time}",
 	}
@@ -45,6 +47,11 @@ func TestNL_P(t *testing.T) {
 		t.Error(err)
 	}
 
+	dur, err := time.ParseDuration("4h2m")
+	if err != nil {
+		t.Error(err)
+	}
+
 	cases := []struct {
 		expression string
 		want       interface{}
@@ -69,6 +76,17 @@ func TestNL_P(t *testing.T) {
 			"time 05-18-1999_6:42pm",
 			tim,
 		},
+		{
+			"dur 4h2m",
+			dur,
+		},
+		{
+			"string Lmao int 42",
+			&T{
+				String: "Lmao",
+				Int:    42,
+			},
+		},
 		{
 			"string What's Up Boy time 05-18-1999_6:42pm",
 			&T{
@@ -101,6 +119,10 @@ func TestNL_P(t *testing.T) {
 				if !c.Time.Equal(v) {
 					t.Errorf("test#%d: got %v want %v", i, c.Time, v)
 				}
+			case time.Duration:
+				if c.Dur != v {
+					t.Errorf("test#%d: got %v want %v", i, c.Dur, v)
+				}
 			case *T:
 				if !reflect.DeepEqual(c, v) {
 					t.Errorf("test#%d: got %v want %v", i, c, v)