alecthomas · brk0v · Jan 15, 2020 · Dorokhov · Jan 17, 2020 · brk0v
diff --git a/bytes.go b/bytes.go
@@ -1,5 +1,9 @@
 package units
 
+import (
+	"errors"
+)
+
 // Base2Bytes is the old non-SI power-of-2 byte scale (1024 bytes in a kilobyte,
 // etc.).
 type Base2Bytes int64
@@ -25,6 +29,8 @@ var (
 	oldBytesUnitMap = MakeUnitMap("B", "B", 1024)
 )
 
+var ErrIncorrectType = errors.New("not a string or int64")
+
 // ParseBase2Bytes supports both iB and B in base-2 multipliers. That is, KB
 // and KiB are both 1024.
 // However "kB", which is the correct SI spelling of 1000 Bytes, is rejected.
@@ -40,6 +46,24 @@ func (b Base2Bytes) String() string {
 	return ToString(int64(b), 1024, "iB", "B")
 }
 
+// UnmarshalTOML implements TOML unmarshal interface.
+func (b *Base2Bytes) UnmarshalTOML(value interface{}) error {
+	switch v := value.(type) {
+	case string:
+		var err error
+		*b, err = ParseBase2Bytes(v)
+		if err != nil {
+			return err
+		}
+	case int64:
+		*b = Base2Bytes(v)
+	default:
+		return ErrIncorrectType
+	}
+
+	return nil
+}
+
 var (
 	metricBytesUnitMap = MakeUnitMap("B", "B", 1000)
 )

diff --git a/bytes_test.go b/bytes_test.go
@@ -99,3 +99,24 @@ func TestParseStrictBytes(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, 1500000, int(n))
 }
+
+func TestUnmarshalTOML(t *testing.T) {
+	b := new(Base2Bytes)
+
+	assert.NoError(t, b.UnmarshalTOML("1KB1B"))
+	assert.Equal(t, 1025, int(*b))
+
+	assert.NoError(t, b.UnmarshalTOML(int64(1000)))
+	assert.Equal(t, 1000, int(*b))
+
+	assert.NoError(t, b.UnmarshalTOML(int64(1001)))
+	assert.Equal(t, 1001, int(*b))
+
+	assert.NoError(t, b.UnmarshalTOML(int64(-10102311)))
+	assert.Equal(t, -10102311, int(*b))
+
+	assert.NoError(t, b.UnmarshalTOML(int64(0)))
+	assert.Equal(t, 0, int(*b))
+
+	assert.Error(t, ErrIncorrectType, b.UnmarshalTOML(uint64(1001)))
+}