Small abstraction to handle picker view elements with a custom object.

PickerView.swift

import Foundation
import ObjectiveC
import UIKit

protocol PickerViewItem: Equatable {

    typealias TextProperty = KeyPath<Self, String>

    typealias SelectProperty = KeyPath<Self, SelectItem>

    associatedtype SelectItem

    var textProperty: TextProperty { get }

    var selectProperty: SelectProperty { get }
}

extension PickerViewItem {

    var text: String {
        return self[keyPath: textProperty]
    }

    var selectedItem: SelectItem {
        return self[keyPath: selectProperty]
    }
}

class PickerView<Element: PickerViewItem>: NSObject, UIPickerViewDataSource, UIPickerViewDelegate {

    typealias Search = (Element) -> Bool

    typealias ToolBarItemsConfiguration = (title: String, color: UIColor)

    typealias SelectedItemEvent = (Element) -> Void

    // MARK: Properties
    private var elements: [Element]

    private var defaultValue: Element?

    private weak var textfield: UITextField?

    // Index of item selected
    private var indexSelected: Int = 0

    var itemSelected: Element?

    var textSelected: String {
        return itemSelected?.text ?? ""
    }

    var selectedValue: Element.SelectItem? {
        return itemSelected?.selectedItem
    }

    private weak var pickerView: UIPickerView?

    // Property holds when a item is selected, also is triggered if a default value was selected
    private var selectedItemEvent: SelectedItemEvent?

    // Property to hold the search function for the default value
    private var search: Search?

    // MARK: Functions
    init(elements: [Element],
         textfield: UITextField,
         pickerView: UIPickerView?,
         defaultWhere search: Search? = nil) {

        self.elements = elements
        self.textfield = textfield
        self.pickerView = pickerView
        self.search = search
        super.init()

        pickerView?.delegate = self
        pickerView?.dataSource = self
        self.textfield?.inputView = pickerView
        self.configureDefaultValue()
    }

    private func configureDefaultValue() {
        guard let search = self.search else {
            return
        }
        defaultValue = self.elements.first(where: search)

        guard let defaultValue = self.defaultValue else {
            return
        }
        self.itemSelected = defaultValue
        selectedItemEvent?(defaultValue)
        setDefaultValueTextField()
    }

    private func setDefaultValueTextField() {
        guard  let defaultValue = self.defaultValue else {
            return
        }

        self.textfield?.text = defaultValue.text
        guard let index = elements.firstIndex(where: {$0 == defaultValue}),
            let pickerView = self.pickerView else {
                return
        }
        self.pickerView?.selectRow(index, inComponent: 0, animated: true)
        self.pickerView?.delegate?.pickerView?(pickerView, didSelectRow: index, inComponent: 0)
    }

    private func resetAll() {
        textfield?.text = nil
        itemSelected = nil
    }

    @objc private func doneButtonPressed() {
        textfield?.resignFirstResponder()

        guard !elements.isEmpty else {
            return
        }

        let element = elements[indexSelected]
        itemSelected = element
        textfield?.text = element.text

        guard let itemSelected = self.itemSelected else {
            return
        }

        selectedItemEvent?(itemSelected)
    }

    @objc private func cancelButtonPressed() {
        textfield?.resignFirstResponder()
    }

    func update(elements: [Element], setDefaultValue: Bool = true) {
        self.elements = elements
        pickerView?.reloadAllComponents()
        if setDefaultValue {
            self.configureDefaultValue()
        }
    }

    @discardableResult
    func addToolBar(tintColor: UIColor,
                    titleConfiguration: ToolBarItemsConfiguration,
                    doneButtonTintColor: UIColor,
                    cancelButton: ToolBarItemsConfiguration,
                    animated: Bool = false) -> Self {

        let toolbar = UIToolbar()
        toolbar.barTintColor = tintColor
        toolbar.sizeToFit()
        let cancelBarButton = UIBarButtonItem(title: cancelButton.title,
                                              style: .plain,
                                              target: self,
                                              action: #selector(cancelButtonPressed))
        cancelBarButton.tintColor = cancelButton.color
        let titleBarButton = UIBarButtonItem(title: titleConfiguration.title, style: .plain, target: nil, action: nil)
        titleBarButton.tintColor = titleConfiguration.color
        let doneBarButton = UIBarButtonItem(barButtonSystemItem: .done, target: self, action: #selector(doneButtonPressed))
        doneBarButton.tintColor = doneButtonTintColor
        let buttons = [cancelBarButton,
                       // This item add more space between buttons
            UIBarButtonItem(barButtonSystemItem: .flexibleSpace, target: nil, action: nil),
            titleBarButton,
            // This item add more space between buttons
            UIBarButtonItem(barButtonSystemItem: .flexibleSpace, target: nil, action: nil),
            doneBarButton]
        toolbar.setItems(buttons, animated: animated)
        textfield?.inputAccessoryView = toolbar

        return self
    }

    @discardableResult
    func selectedItem(_ selectedItemEvent: @escaping SelectedItemEvent) -> Self {
        self.selectedItemEvent = selectedItemEvent
        return self
    }

    // MARK: DataSource/Delegate
    func numberOfComponents(in pickerView: UIPickerView) -> Int {
        return 1
    }

    func pickerView(_ pickerView: UIPickerView, numberOfRowsInComponent component: Int) -> Int {
        return elements.count
    }

    func pickerView(_ pickerView: UIPickerView, didSelectRow row: Int, inComponent component: Int) {
        indexSelected = row
    }

    func pickerView(_ pickerView: UIPickerView, titleForRow row: Int, forComponent component: Int) -> String? {
        let element =  elements[row]
        return element.text
    }
}

var associatedObjectHandle: UInt8 = 0

extension UITextField {

    func pickerView<T: PickerViewItem>(itemType type: T.Type = T.self) -> PickerView<T>? {
        return objc_getAssociatedObject(self, &associatedObjectHandle) as? PickerView<T>
    }

    func addPickerView<T: PickerViewItem>(_ pickerView: UIPickerView = .init(),
                                          elements: [T],
                                          defaultWhere predicate: PickerView<T>.Search? = nil) -> PickerView<T>? {
        // Prevent reset a pre set element
        if objc_getAssociatedObject(self, &associatedObjectHandle) == nil {
            let pickerView = PickerView(elements: elements, textfield: self, pickerView: pickerView, defaultWhere: predicate)
            objc_setAssociatedObject(self, &associatedObjectHandle, pickerView, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
            return pickerView
        } else {
            return self.pickerView(itemType: T.self)
        }
    }
}


// USAGE

struct Document: PickerViewItem {
    let name: String
    let code: String

    var selectProperty: KeyPath<Document, String> {
        return \.code
    }

    var textProperty: KeyPath<Document, String> {
        return \.name
    }
}

let textfield = UITextField()

// creating picker
textfield.addPickerView(elements: [ Document(name: "License", code: "D"),
                                    Document(name: "Id Number", code: "I"),
                                    Document(name: "Passport", code: "P")])?
    // adding toolbar if you needed
    .addToolBar(tintColor: .black, titleConfiguration: (title: "Title", color: .red), doneButtonTintColor: .red, cancelButton: (title: "Cancel", color: .red))
    // get object when a item is selected
    .selectedItem { document in
        print("document \(document)")
    }

// updating elements
textfield.pickerView(itemType: Document.self)?.update(elements: [Document(name: "Passport", code: "P")])

// getting elements

let selectedValue = textfield.pickerView(itemType: Document.self)?.selectedValue
let textSelected = textfield.pickerView(itemType: Document.self)?.textSelected
let itemSelected = textfield.pickerView(itemType: Document.self)?.itemSelected