Custom UnmarshalYAML interface for an interface and its implementations

问题

I implemented an interface Fruit and two implementations of it: Apple and Banana.

Into objects of the two implementations I want to load data from a yaml file:

capacity: 4
Apple:
- name: "apple1"
  number: 1
- name: "apple2"
  number: 1
Banana:
- name: "banana1"
  number: 2

I implemented the UnmarshalYaml interface to load data into my objects:

package main

import (
    "errors"
    "gopkg.in/yaml.v3"
    "log"
    "fmt"
)

type FruitBasket struct {
    Capacity int `yaml:"capacity"`
    Fruits []Fruit
}

func NewFruitBasket() *FruitBasket {
    fb := new(FruitBasket)

    return fb
}

type Fruit interface {
    GetFruitName() string
    GetNumber() int
}

type Apple struct {
    Name string `yaml:"name"`
    Number int `yaml:"number"`
}

type Banana struct {
    Name string `yaml:"name"`
    Number int `yaml:"number"`
}

func (apple *Apple) GetFruitName() string {
    return apple.Name
}

func (apple *Apple) GetNumber() int {
    return apple.Number
}

func (banana *Banana) GetFruitName() string {
    return banana.Name
}

func (banana *Banana) GetNumber() int {
    return banana.Number
}

type tmpFruitBasket struct {
    Capacity int `yaml:"capacity"`
    Fruits []map[string]yaml.Node
}

func (fruitBasket *FruitBasket) UnmarshalYAML(value *yaml.Node) error {
    var tmpFruitBasket tmpFruitBasket

    if err := value.Decode(&tmpFruitBasket); err != nil {
        return err
    }

    fruitBasket.Capacity = tmpFruitBasket.Capacity

    fruits := make([]Fruit, 0, len(tmpFruitBasket.Fruits))

    for i := 0; i < len(tmpFruitBasket.Fruits); i++ {
        for tag, node := range tmpFruitBasket.Fruits[i] {
            switch tag {
            case "Apple":
                apple := &Apple{}
                if err := node.Decode(apple); err != nil {
                    return err
                }

                fruits = append(fruits, apple)
            case "Banana":
                banana := &Banana{}
                if err := node.Decode(banana); err != nil {
                    return err
                }

                fruits = append(fruits, banana)
            default:
                return errors.New("Failed to interpret the fruit of type: \"" + tag + "\"")
            }
        }
    }

    fruitBasket.Fruits = fruits

    return nil
}

func main() {
    data := []byte(`
capacity: 4
Apple:
- name: "apple1"
  number: 1
- name: "apple2"
  number: 1
Banana:
- name: "banana1"
  number: 2
`)

    fruitBasket := NewFruitBasket()

    err := yaml.Unmarshal(data, &fruitBasket)

    if err != nil {
        log.Fatalf("error: %v", err)
    }

    fmt.Println(fruitBasket.Capacity)

    for i := 0; i < len(fruitBasket.Fruits); i++ {
        switch fruit := fruitBasket.Fruits[i].(type) {
        case *Apple:
            fmt.Println(fruit.Name)
            fmt.Println(fruit.Number)
        }
    }
}

However, this is not working. It seems that the data for the Apple and Banana tags are not loaded. Probably, because of the missing yaml flag for the Fruits slice in my tmpFruitBasket struct. But, as Fruit is an interface, I cannot define a yaml flag. In the future, I want to implement other structs representing concrete fruits (e.g., Strawberry) implementing the interface Fruit.

Any idea on how to solve this?

回答1:

This is the intermediate type you need:

type tmpFruitBasket struct {
  Capacity int
  Apple    []yaml.Node `yaml:"Apple"`
  Banana   []yaml.Node `yaml:"Banana"`
}

Then, the loading function will look like:

// helper to load a list of nodes as a concrete type
func appendFruits(fruits []Fruit, kind reflect.Type, input []yaml.Node) ([]Fruit, error) {
  for i := range input {
    val := reflect.New(kind).Interface()
    if err := input[i].Decode(val); err != nil {
      return nil, err
    }
    fruits = append(fruits, val.(Fruit))
  }
  return fruits, nil
}


func (fruitBasket *FruitBasket) UnmarshalYAML(value *yaml.Node) error {
    var tmp tmpFruitBasket

    if err := value.Decode(&tmp); err != nil {
        return err
    }

    fruitBasket.Capacity = tmp.Capacity

    var fruits []Fruit
    var err error
    // sadly, there is no nicer way to get the reflect.Type of Apple / Banana
    fruits, err = appendFruits(
      fruits, reflect.TypeOf((*Apple)(nil)).Elem(), tmp.Apple)
    if err != nil {
      return err
    }
    fruits, err = appendFruits(
      fruits, reflect.TypeOf((*Banana)(nil)).Elem(), tmp.Banana)
    if err != nil {
      return err
    }

    fruitBasket.Fruits = fruits
    return nil
}

Edit: If you stick with sorting each type into a dedicated slice, you could of course directly type those as []Apple and []Banana and just merge them. This answer is a continuation of digging into the issue of dynamically loading input to different types, started with your previous questions. Doing this only makes sense if you at some point don't know the static type anymore (e.g. if you provide an API to add additional fruit types at runtime).

来源：https://stackoverflow.com/questions/60636689/custom-unmarshalyaml-interface-for-an-interface-and-its-implementations