Three ways of registering a callback and using keypaths to avoid forgetting to use [weak self]

Playground.swift

protocol SafeCallbackProtocol {
    func invoke(message: String)
}

struct SafeCallback<Root: AnyObject>: SafeCallbackProtocol {
    weak var target: Root?
    let path: ReferenceWritableKeyPath<Root, String>

    init(target: Root, path: ReferenceWritableKeyPath<Root, String>) {
        self.target = target
        self.path = path
    }

    func invoke(message: String) {
        target?[keyPath: path] = message
    }
}

class ModelA {
    let modelB: ModelB

    private var message: String {
        get { "" }
        set { printMessage(newValue) }
    }

    init(modelB: ModelB) {
        self.modelB = modelB
    }

    private func printMessage(_ message: String) {
        print("ModelA received: \(message)")
    }

    public func registerCallback1() {
        modelB.registerCallback { message in
            self.printMessage(message)
        }
    }

    public func registerCallback2() {
        modelB.registerCallback { [weak self] message in
            self?.printMessage(message)
        }
    }

    public func registerCallbackSafely() {
        modelB.registerSafeCallback(to: self, \.message)
    }
}

class ModelB {
    var callback: ((String) -> Void)?
    var safeCallback: SafeCallbackProtocol?

    func registerCallback(_ callback: @escaping (String) -> Void) {
        self.callback = callback
    }

    func registerSafeCallback<Root: AnyObject>(to root: Root, _ path: ReferenceWritableKeyPath<Root, String>) {
        safeCallback = SafeCallback(target: root, path: path)
    }

    func invokeCallback(_ message: String) {
        callback?(message)
        safeCallback?.invoke(message: message)
    }
}

var strongModel1: ModelA? = ModelA(modelB: ModelB())
weak var weakModel1: ModelA? = strongModel1
strongModel1?.registerCallback1()
strongModel1?.modelB.invokeCallback("Hello 1")

var strongModel2: ModelA? = ModelA(modelB: ModelB())
weak var weakModel2: ModelA? = strongModel2
strongModel2?.registerCallback2()
strongModel2?.modelB.invokeCallback("Hello 2")

var strongModel3: ModelA? = ModelA(modelB: ModelB())
weak var weakModel3: ModelA? = strongModel3
strongModel3?.registerCallbackSafely()
strongModel3?.modelB.invokeCallback("Hello 3")

strongModel1 = nil
strongModel2 = nil
strongModel3 = nil
print("weakModel1 = \(weakModel1)")
print("weakModel2 = \(weakModel2)")
print("weakModel3 = \(weakModel3)")

/* Output is:

 ModelA received: Hello 1
 ModelA received: Hello 2
 ModelA received: Hello 3
 weakModel1 = Optional(__lldb_expr_19.ModelA)
 weakModel2 = nil
 weakModel3 = nil

 */

Updated to show more realistic API

    func registerSafeCallback<Root: AnyObject>(to root: Root, _ path: ReferenceWritableKeyPath<Root, String>) {
        safeCallback = SafeCallback(target: root, path: path)
    }

Usage:

modelB.registerSafeCallback(to: self, \.message)