Skip to content

Instantly share code, notes, and snippets.

@luw2007
Last active November 28, 2019 08:37
Show Gist options
  • Save luw2007/172c72b3cab34f765e032fff469ef316 to your computer and use it in GitHub Desktop.
Save luw2007/172c72b3cab34f765e032fff469ef316 to your computer and use it in GitHub Desktop.
AddMonths 增加月份,有许多神奇的case需要处理,例如 1月31日 - 2月28日, 3月31日 - 4月30日。 总的逻辑为,到达下一个月的同一天,若无效前推至有效
package main
import (
"fmt"
"testing"
"time"
)
// WHY:
//
// NextMonths plus or minus natural month; many magical cases, such as 0131 - 0228, 0331 - 0430.
// reach the same day of the next month, reset the last month if it is invalid
func NextMonths(t time.Time, n int) time.Time {
y, m, d := t.Date()
hour, min, sec := t.Clock()
next := time.Date(y, m+time.Month(n), d, hour, min, sec, 0, t.Location())
if (next.Year()-y)*12+int(next.Month()-m) == n {
return next
}
return time.Date(y, m+time.Month(n+1), 0, hour, min, sec, 0, t.Location())
}
func mustParseTime(s string) time.Time {
t, _ := time.Parse("20060102", s)
return t
}
func TestNextMonths(t *testing.T) {
type args struct {
start time.Time
n int
}
tests := []struct {
name string
args args
want time.Time
}{
{"0130-0228", args{mustParseTime("20190130"), 1}, mustParseTime("20190228")},
{"0131-0228", args{mustParseTime("20190131"), 1}, mustParseTime("20190228")},
{"1231-0229", args{mustParseTime("20191231"), 2}, mustParseTime("20200229")},
{"0101-0101", args{mustParseTime("20190101"), 0}, mustParseTime("20190101")},
{"0101-0201", args{mustParseTime("20190101"), 1}, mustParseTime("20190201")},
{"0331-0430", args{mustParseTime("20190331"), 1}, mustParseTime("20190430")},
{"0731-0831", args{mustParseTime("20190731"), 13}, mustParseTime("20200831")},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := NextMonths(tt.args.start, tt.args.n); got != tt.want {
t.Errorf("NextMonths() = %v, want %v", got, tt.want)
}
})
}
}
func main() {
// 20190131 -> 20190228
fmt.Println(NextMonths(mustParseTime("20190130"), 1))
// 20190331 -> 20190430
fmt.Println(NextMonths(mustParseTime("20190331"), 1))
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment