best implementation of period and marshalling

aggregator.go

@@ -29,7 +29,7 @@ func NewTimeAggregator(units ...unit) (*TimeAggregator, error) {
 }
 
 func (a *TimeAggregator) Add(date time.Time, value int64) {
-	p := pack(a.flags, date)
+	p := NewPeriod(a.flags, date)
 	if _, ok := a.Values[p]; !ok {
 		a.Values[p] = a.buildAggregator()
 	}
@@ -38,7 +38,7 @@ func (a *TimeAggregator) Add(date time.Time, value int64) {
 }
 
 func (a *TimeAggregator) Get(date time.Time) int64 {
-	p := pack(a.flags, date)
+	p := NewPeriod(a.flags, date)
 
 	if _, ok := a.Values[p]; !ok {
 		return -1
@@ -63,8 +63,9 @@ func (a *TimeAggregator) Marshal() []byte {
 }
 
 func (a *TimeAggregator) Unmarshal(v []byte) error {
-	buf := bytes.NewBuffer(v)
+	a.Values = make(map[Period]*aggregator, 0)
 
+	buf := bytes.NewBuffer(v)
 	for {
 		var p Period
 		if err := binary.Read(buf, binary.LittleEndian, &p); err != nil {
@@ -75,6 +76,12 @@ func (a *TimeAggregator) Unmarshal(v []byte) error {
 			return err
 		}
 
+		if a.flags == 0 {
+			a.flags = p.Flag()
+			us := p.Units()
+			a.kind = us[len(us)-1]
+		}
+
 		a.Values[p] = a.buildAggregator()
 		if err := a.Values[p].Unmarshal(buf); err != nil {
 			return err

aggregator_test.go

@@ -18,25 +18,6 @@ func (s *UtilsSuite) TestNewTimeAggregator(c *C) {
 	c.Assert(err, Equals, InvalidOrderError)
 }
 
-func (s *UtilsSuite) TestTimeAggregator_Add(c *C) {
-	d := time.Now()
-
-	a, _ := NewTimeAggregator(Year, Hour)
-	a.Add(d, 10)
-	a.Add(d, 10)
-
-	c.Assert(a.Values, HasLen, 1)
-	c.Assert(a.Get(d), Equals, int64(20))
-
-	m := a.Marshal()
-
-	b, _ := NewTimeAggregator(Year, Hour)
-	err := b.Unmarshal(m)
-	c.Assert(err, IsNil)
-	c.Assert(b.Values, HasLen, 1)
-	c.Assert(b.Get(d), Equals, int64(20))
-}
-
 func (s *UtilsSuite) TestTimeAggregator_Add_YearHour(c *C) {
 	a, _ := NewTimeAggregator(Year, Hour)
 	a.Add(date2014November, 15)
@@ -84,6 +65,27 @@ func (s *UtilsSuite) TestTimeAggregator_Add_Only(c *C) {
 	c.Assert(a.Get(h21), Equals, int64(40))
 }
 
+func (s *UtilsSuite) TestTimeAggregator_MarshalAndUnmarshal(c *C) {
+	d := time.Now()
+
+	a, _ := NewTimeAggregator(Year, Hour)
+	a.Add(d, 10)
+	a.Add(d, 10)
+
+	c.Assert(a.Values, HasLen, 1)
+	c.Assert(a.Get(d), Equals, int64(20))
+
+	m := a.Marshal()
+
+	b := &TimeAggregator{}
+	err := b.Unmarshal(m)
+
+	c.Assert(err, IsNil)
+	c.Assert(a.flags, Equals, b.flags)
+	c.Assert(b.Values, HasLen, 1)
+	c.Assert(b.Get(d), Equals, int64(20))
+}
+
 var date2013December = time.Date(2013, time.December, 12, 23, 59, 59, 0, time.UTC)
 var date2013November = time.Date(2013, time.November, 12, 23, 59, 59, 0, time.UTC)
 var date2014February = time.Date(2014, time.February, 12, 23, 59, 59, 0, time.UTC)

periods.go

@@ -4,8 +4,6 @@ import (
 	"time"
 )
 
-type Period uint64
-
 type unit int
 
 const (
@@ -106,7 +104,9 @@ var defs = map[unit]periodDefinition{
 	},
 }
 
-func pack(flag unit, date time.Time) Period {
+type Period uint64
+
+func NewPeriod(flag unit, date time.Time) Period {
 	us := getUnitsFromFlag(flag)
 
 	t := binaryVersion
@@ -121,10 +121,10 @@ func pack(flag unit, date time.Time) Period {
 	return Period(t)
 }
 
-func unpack(total Period) map[string]uint64 {
-	t := uint64(total)
+func (p Period) ToMap() map[string]uint64 {
+	t := uint64(p)
 
-	us := getUnitsFromFlag(unit(t % 1e3))
+	us := p.Units()
 	t = t / 1e3
 
 	result := make(map[string]uint64, 0)
@@ -141,6 +141,14 @@ func unpack(total Period) map[string]uint64 {
 	return result
 }
 
+func (p Period) Flag() unit {
+	return unit(uint64(p) % 1e3)
+}
+
+func (p Period) Units() []unit {
+	return getUnitsFromFlag(p.Flag())
+}
+
 func getUnitsFromFlag(flag unit) []unit {
 	us := make([]unit, 0)
 	for _, u := range units {

periods_test.go

@@ -4,12 +4,12 @@ import (
 	. "gopkg.in/check.v1"
 )
 
-func (s *UtilsSuite) Test_pack(c *C) {
-	flags := Year | Month | Hour
-	c.Assert(pack(flags, date2015November), Equals, Period(1201511067))
+func (s *UtilsSuite) Test_NewPeriod(c *C) {
+	p := NewPeriod(Year|Month|Hour, date2015November)
+	c.Assert(p, Equals, Period(1201511067))
 }
 
 func (s *UtilsSuite) Test_unpack(c *C) {
-	p := unpack(Period(1201511067))
-	c.Assert(p, DeepEquals, map[string]uint64{"year": 2015, "month": 11})
+	p := Period(1201511067)
+	c.Assert(p.ToMap(), DeepEquals, map[string]uint64{"year": 2015, "month": 11})
 }