From b6fe702588a6e9053045bbcce86b9193b65c0d65 Mon Sep 17 00:00:00 2001
From: Martin Tournoij <martin@arp242.net>
Date: Sun, 1 Oct 2023 10:25:02 +0100
Subject: [PATCH] Wrap UnmarshalTOML()/UnmarshalText return values in
 ParseError

This adds position information.

Fixes #398
---
 decode.go     |  8 +++--
 error_test.go | 85 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 2 deletions(-)

diff --git a/decode.go b/decode.go
index 48a8d405..210bd3a8 100644
--- a/decode.go
+++ b/decode.go
@@ -218,7 +218,11 @@ func (md *MetaData) unify(data any, rv reflect.Value) error {
 
 	rvi := rv.Interface()
 	if v, ok := rvi.(Unmarshaler); ok {
-		return v.UnmarshalTOML(data)
+		err := v.UnmarshalTOML(data)
+		if err != nil {
+			return md.parseErr(err)
+		}
+		return nil
 	}
 	if v, ok := rvi.(encoding.TextUnmarshaler); ok {
 		return md.unifyText(data, v)
@@ -533,7 +537,7 @@ func (md *MetaData) unifyText(data any, v encoding.TextUnmarshaler) error {
 		return md.badtype("primitive (string-like)", data)
 	}
 	if err := v.UnmarshalText([]byte(s)); err != nil {
-		return err
+		return md.parseErr(err)
 	}
 	return nil
 }
diff --git a/error_test.go b/error_test.go
index 8da5d9d7..a7de72c0 100644
--- a/error_test.go
+++ b/error_test.go
@@ -266,3 +266,88 @@ func TestParseError(t *testing.T) {
 		})
 	}
 }
+
+type Enum1 uint8
+
+func (n *Enum1) UnmarshalText(text []byte) error {
+	switch t := strings.TrimSpace(string(text)); t {
+	case "ok":
+		*n = 1
+	default:
+		return fmt.Errorf("invalid value: %q", t)
+	}
+	return nil
+}
+
+// Make sure custom types are wrapped in ParseError with correct location.
+func TestUnmarshalTypeError(t *testing.T) {
+	var c struct {
+		K1 string `toml:"k1"`
+		K2 Enum1  `toml:"k2"`
+		K3 Enum1  `toml:"k3"`
+	}
+	_, err := toml.Decode("k1 = 'asd'\nk2 = 'ok'\nk3 = 'invalid'\nk4 = 'ok'", &c)
+	if err == nil {
+		t.Fatal("error is nil")
+	}
+	var pErr toml.ParseError
+	if !errors.As(err, &pErr) {
+		t.Fatalf("not a ParseError: %#v", err)
+	}
+
+	want := `toml: error: invalid value: "invalid"
+
+At line 3, column 6-13:
+
+      1 | k1 = 'asd'
+      2 | k2 = 'ok'
+      3 | k3 = 'invalid'
+                ^^^^^^^
+`
+
+	if have := pErr.ErrorWithUsage(); have != want {
+		t.Errorf("\nwant:\n%s\nhave:\n%s", want, have)
+	}
+}
+
+type Enum2 uint8
+
+func (n *Enum2) UnmarshalTOML(text any) error {
+	switch t := strings.TrimSpace(text.(string)); t {
+	case "ok":
+		*n = 1
+	default:
+		return fmt.Errorf("invalid value: %q", t)
+	}
+	return nil
+}
+
+func TestMarhsalError(t *testing.T) {
+	var c struct {
+		K1 string `toml:"k1"`
+		K2 Enum2  `toml:"k2"`
+		K3 Enum2  `toml:"k3"`
+	}
+	_, err := toml.Decode("k1 = 'asd'\nk2 = 'ok'\nk3 = 'invalid'\nk4 = 'ok'", &c)
+	if err == nil {
+		t.Fatal("error is nil")
+	}
+	var pErr toml.ParseError
+	if !errors.As(err, &pErr) {
+		t.Fatalf("not a ParseError: %#v", err)
+	}
+
+	want := `toml: error: invalid value: "invalid"
+
+At line 3, column 6-13:
+
+      1 | k1 = 'asd'
+      2 | k2 = 'ok'
+      3 | k3 = 'invalid'
+                ^^^^^^^
+`
+
+	if have := pErr.ErrorWithUsage(); have != want {
+		t.Errorf("\nwant:\n%s\nhave:\n%s", want, have)
+	}
+}