GolangでEnumをフィールドに持つstructをいい感じにjsonエンコード / デコードする

最近、Golang で Enum をフィールドに持つ struct をいい感じに json エンコード / デコードしたい衝動に駆られた事があったので、その時のメモです。

普通に json.Marshal する

例えば、以下のような Userの struct を定義したとする。

type UserRole int

const (
    _ UserRole = iota
    Consumer
    Admin
)

func (r UserRole) String() string {
    switch r {
    case Consumer:
        return "consumer"
    case Admin:
        return "admin"
    default:
        return "unknown"
    }
}

type User struct {
    Id   string   `json:"id"`
    Role UserRole `json:"role"`
}

こいつを普通にjson.Marshalしてみる。

func main() {
    u := &User{Id: "abcde", Role: Admin}
    b, err := json.Marshal(u)
    if err != nil {
        log.Fatal(err)
    }
    fmt.Println(string(b)) // {"id":"abcde","role":2}
}

これが{"id":"abcde","role":"admin"}のような json になってほしかったのだけど、ただjson.Marshalしただけでは{"id":"abcde","role":2}という文字列が得られてツラい思いをした。

Enum の MarshalJSON メソッドを定義する

type Marshaler interface {
        MarshalJSON() ([]byte, error)
}

というinterfaceがあって、MarshalJSONを実装しておけばいい感じに json エンコードしてくれそうだということがわかったので試してみる。

// 上述のコードにこれを追加
func (r UserRole) MarshalJSON() ([]byte, error) {
    return json.Marshal(r.String())
}

今度はちゃんと{"id":"abcde","role":"admin"}が得られた。

ただしMarshalJSONを実装しただけでは、json 文字列を再度 User にしたい時に、 json: cannot unmarshal string into Go value of type main.UserRoleのようなエラーが出てまたツラい思いをする。

func main() {
    u := &User{Id: "abcde", Role: Admin}
    b, err := json.Marshal(u)
    if err != nil {
        log.Fatal(err)
    }
    fmt.Println(string(b)) // {"id":"abcde","role":2}

    var u2 User
    err = json.Unmarshal(b, &u2)
    if err != nil {
        log.Fatal(err) // json: cannot unmarshal string into Go value of type main.UserRole
    }
    fmt.Println(u2)
}

これに関してはMarshalJSONの時と同様にUnmarshalJSONを実装しておけばよさそう。

func (r *UserRole) UnmarshalJSON(data []byte) error {
    var s string
    if err := json.Unmarshal(data, &s); err != nil {
        return fmt.Errorf("data should be a string, got %s", data)
    }

    var ur UserRole
    switch s {
    case "consumer":
        ur = Consumer
    case "admin":
        ur = Admin
    default:
        return fmt.Errorf("invalid UserRole %s", s)

    }
    *r = ur
    return nil
}

これで、ちゃんと json から{abcde admin}というUserの struct が得られる。

jsonenums を使う

最初は上記のように頑張って自分でMarshalJSON, UnmarshalJSONを定義していたが、もう少し調べてみたらjsonenumsというライブラリを使えばこれらの作業を自動化してくれるっぽい。というわけで最終的にこれを使い倒す形で落ち着いた。

プロジェクト構成を以下のようにしたとすると、modelsディレクトリ以下でjsonenums -type=UserRoleとするだけでuserrole_jsonenums.goを自動生成してくれる。

.
├── main.go
└── models
    ├── models.go
    └── userrole_jsonenums.go // jsonenumsが生成してくれる

最終的にコードは以下のようになった。

// models.go

package models

type UserRole int

const (
    _ UserRole = iota
    Consumer
    Admin
)

func (r UserRole) String() string {
    switch r {
    case Consumer:
        return "consumer"
    case Admin:
        return "admin"
    default:
        return "unknown"
    }
}

type User struct {
    Id   string   `json:"id"`
    Role UserRole `json:"role"`
}

// userrole_jsonenums.go

// generated by jsonenums -type=UserRole; DO NOT EDIT

package models

import (
    "encoding/json"
    "fmt"
)

var (
    _UserRoleNameToValue = map[string]UserRole{
        "Consumer": Consumer,
        "Admin":    Admin,
    }

    _UserRoleValueToName = map[UserRole]string{
        Consumer: "Consumer",
        Admin:    "Admin",
    }
)

func init() {
    var v UserRole
    if _, ok := interface{}(v).(fmt.Stringer); ok {
        _UserRoleNameToValue = map[string]UserRole{
            interface{}(Consumer).(fmt.Stringer).String(): Consumer,
            interface{}(Admin).(fmt.Stringer).String():    Admin,
        }
    }
}

// MarshalJSON is generated so UserRole satisfies json.Marshaler.
func (r UserRole) MarshalJSON() ([]byte, error) {
    if s, ok := interface{}(r).(fmt.Stringer); ok {
        return json.Marshal(s.String())
    }
    s, ok := _UserRoleValueToName[r]
    if !ok {
        return nil, fmt.Errorf("invalid UserRole: %d", r)
    }
    return json.Marshal(s)
}

// UnmarshalJSON is generated so UserRole satisfies json.Unmarshaler.
func (r *UserRole) UnmarshalJSON(data []byte) error {
    var s string
    if err := json.Unmarshal(data, &s); err != nil {
        return fmt.Errorf("UserRole should be a string, got %s", data)
    }
    v, ok := _UserRoleNameToValue[s]
    if !ok {
        return fmt.Errorf("invalid UserRole %q", s)
    }
    *r = v
    return nil
}

// main.go

package main

import (
    "encoding/json"
    "fmt"
    "log"

    "github.com/danimal141/enum-test/models"
)

func main() {
    u := &models.User{Id: "abcde", Role: models.Admin}
    b, err := json.Marshal(u)
    if err != nil {
        log.Fatal(err)
    }
    fmt.Println(string(b)) // {"id":"abcde","role":"admin"} よい

    var u2 models.User
    err = json.Unmarshal(b, &u2)
    if err != nil {
        log.Fatal(err)
    }
    fmt.Println(u2) // {abcde admin} よい
}

まとめ

Enum をフィールドに持つ struct の json エンコード / デコードはjsonenumsを使うか、自分でMarshalJSON, UnmarshalJSONを定義すればよさそう。 何か他によい方法がありそうでしたらぜひ教えてください。

参考

• jsonenums