From 000d30f9a02a4b5436758d2eccefebfb808fc0e8 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Fri, 27 Sep 2024 20:42:24 +0300
Subject: [PATCH] exfmt: make duration formatting more customizable

---
 exfmt/duration.go      | 58 ++++++++++++++++++++++++++++++++----------
 exfmt/duration_test.go |  6 +++++
 2 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/exfmt/duration.go b/exfmt/duration.go
index b12ca68..c95ce49 100644
--- a/exfmt/duration.go
+++ b/exfmt/duration.go
@@ -16,35 +16,67 @@ import (
 var Day = 24 * time.Hour
 var Week = 7 * Day
 
-func pluralize(value int, unit string) string {
-	if value == 1 {
-		return "1 " + unit
+type Pluralizer func(int) string
+
+func Pluralizable(unit string) Pluralizer {
+	return func(value int) string {
+		if value == 1 {
+			return "1 " + unit
+		}
+		return fmt.Sprintf("%d %ss", value, unit)
+	}
+}
+
+func NonPluralizable(unit string) Pluralizer {
+	return func(value int) string {
+		return fmt.Sprintf("%d %s", value, unit)
 	}
-	return fmt.Sprintf("%d %ss", value, unit)
 }
 
-func appendDurationPart(time, unit time.Duration, name string, parts *[]string) (remainder time.Duration) {
+func Duration(d time.Duration) string {
+	return DurationCustom(d, nil, Week, Day, time.Hour, time.Minute, time.Second)
+}
+
+func appendDurationPart(time, unit time.Duration, name Pluralizer, parts *[]string) (remainder time.Duration) {
 	if time < unit {
 		return time
 	}
 	value := int(time / unit)
 	remainder = time % unit
-	*parts = append(*parts, pluralize(value, name))
+	*parts = append(*parts, name(value))
 	return
 }
 
-func Duration(d time.Duration) string {
+var DefaultDurationUnitNames = map[time.Duration]Pluralizer{
+	Week:             Pluralizable("week"),
+	Day:              Pluralizable("day"),
+	time.Hour:        Pluralizable("hour"),
+	time.Minute:      Pluralizable("minute"),
+	time.Second:      Pluralizable("second"),
+	time.Millisecond: NonPluralizable("ms"),
+	time.Microsecond: NonPluralizable("µs"),
+	time.Nanosecond:  NonPluralizable("ns"),
+}
+
+func DurationCustom(d time.Duration, names map[time.Duration]Pluralizer, units ...time.Duration) string {
 	if d < 0 {
 		panic(errors.New("exfmt.Duration: negative duration"))
-	} else if d < time.Second {
+	} else if len(units) == 0 {
+		panic(errors.New("exfmt.Duration: no units provided"))
+	} else if d < units[len(units)-1] {
 		return "now"
 	}
+	if names == nil {
+		names = DefaultDurationUnitNames
+	}
 	parts := make([]string, 0, 2)
-	d = appendDurationPart(d, Week, "week", &parts)
-	d = appendDurationPart(d, Day, "day", &parts)
-	d = appendDurationPart(d, time.Hour, "hour", &parts)
-	d = appendDurationPart(d, time.Minute, "minute", &parts)
-	_ = appendDurationPart(d, time.Second, "second", &parts)
+	for _, unit := range units {
+		name, ok := names[unit]
+		if !ok {
+			panic(fmt.Errorf("exfmt.Duration: no name for unit %q", unit))
+		}
+		d = appendDurationPart(d, unit, name, &parts)
+	}
 	if len(parts) > 2 {
 		parts[0] = strings.Join(parts[:len(parts)-1], ", ")
 		parts[1] = parts[len(parts)-1]
diff --git a/exfmt/duration_test.go b/exfmt/duration_test.go
index c6c4c47..946dc72 100644
--- a/exfmt/duration_test.go
+++ b/exfmt/duration_test.go
@@ -59,3 +59,9 @@ func TestFormatDuration_PanicNegative(t *testing.T) {
 	assert.Panics(t, func() { exfmt.Duration(-time.Second) })
 	assert.Panics(t, func() { exfmt.Duration(-exfmt.Week) })
 }
+
+func TestFormatDuration_Custom(t *testing.T) {
+	assert.Equal(t, "90 days", exfmt.DurationCustom(90*exfmt.Day, nil, exfmt.Day))
+	assert.Equal(t, "2160 hours", exfmt.DurationCustom(90*exfmt.Day, nil, time.Hour))
+	assert.Equal(t, "2160 hours", exfmt.DurationCustom(90*exfmt.Day+59*time.Minute, nil, time.Hour))
+}