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)