Decoding heterogeneous array with Swift Decodable

HeterogeneousListDecoding.swift

import Foundation

enum TypeKey: CodingKey {
    case type
}

protocol TypeFamily: Decodable {
    associatedtype Member: Decodable
    var type: Member.Type { get }
}

protocol KeyedListKey: CodingKey {
    static var key: Self { get }
}

struct UnkeyedHeterogeneousList<Family: TypeFamily>: Decodable {
    let list: [Family.Member]

    init(from decoder: Decoder) throws {
        var typeContainer = try decoder.unkeyedContainer()
        var list: [Family.Member] = []

        var listContainer = typeContainer
        while !listContainer.isAtEnd {
            let instance = try typeContainer.nestedContainer(keyedBy: TypeKey.self)
            if let type = try? instance.decode(Family.self, forKey: .type) {
                list.append(try listContainer.decode(type.type))
            } else {
                list.append(try listContainer.decode(Family.Member.self))
            }
        }

        self.list = list
    }
}

struct KeyedHeterogeneousList<Family: TypeFamily, Key: KeyedListKey>: Decodable {
    let list: [Family.Member]

    init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: Key.self)
        var typeContainer = try container.nestedUnkeyedContainer(forKey: Key.key)
        var list: [Family.Member] = []

        var listContainer = typeContainer
        while !listContainer.isAtEnd {
            let instance = try typeContainer.nestedContainer(keyedBy: TypeKey.self)
            if let type = try? instance.decode(Family.self, forKey: .type) {
                list.append(try listContainer.decode(type.type))
            } else {
                list.append(try listContainer.decode(Family.Member.self))
            }
        }

        self.list = list
    }
}

HeterogeneousListDecodingTests.swift

import XCTest

private class Shape: Decodable {
    let id: Int
}

private class Square: Shape {
    let length: Double

    private enum CodingKeys: String, CodingKey {
        case length
    }

    required init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        length = try container.decode(Double.self, forKey: .length)

        try super.init(from: decoder)
    }
}

private class Circle: Shape {
    let radius: Double

    private enum CodingKeys: String, CodingKey {
        case radius
    }

    required init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        radius = try container.decode(Double.self, forKey: .radius)

        try super.init(from: decoder)
    }
}

private enum ShapeFamily: String, TypeFamily {
    case square
    case circle

    var type: Shape.Type {
        switch self {
        case .square:
            return Square.self
        case .circle:
            return Circle.self
        }
    }
}

private enum ShapeListKey: KeyedListKey {
    case shapes

    static var key = Self.shapes
}

final class HeterogeneousDecodingTests: XCTestCase {
    var unkeyedData: String {
        """
        [
            {
                "id": 1,
                "type": "square",
                "length": 10
            },
            {
                "id": 2,
                "type": "circle",
                "radius": 3.14
            },
            {
                "id": 3,
                "type": "unknown"
            }
        ]
        """
    }

    var keyedData: String {
        """
        {
            "shapes": \(unkeyedData)
        }
        """
    }

    func testUnkeyedDecoding() {
        let result = try! JSONDecoder().decode(
            UnkeyedHeterogeneousList<ShapeFamily>.self,
            from: Data(unkeyedData.utf8)
        )

        checkDecoded(shapes: result.list)
    }

    func testKeyedDecoding() {
        let result = try! JSONDecoder().decode(
            KeyedHeterogeneousList<ShapeFamily, ShapeListKey>.self,
            from: Data(keyedData.utf8)
        )

        checkDecoded(shapes: result.list)
    }

    private func checkDecoded(shapes: [Shape]) {
        let square = shapes[0] as! Square
        XCTAssertEqual(square.id, 1)
        XCTAssertEqual(square.length, 10)

        let circle = shapes[1] as! Circle
        XCTAssertEqual(circle.id, 2)
        XCTAssertEqual(circle.radius, 3.14)

        let unknown = shapes[2]
        XCTAssertEqual(unknown.id, 3)
    }
}