Custom trait propagation not working as expected outside of view hierarchy

Question

I'm trying to build custom-styled components that can handle a custom trait propagated from above and adjust their styling accordingly. E.g.:

class CustomButton: UIButton {
    override init(frame: CGRect) {
      super.init(frame: frame)

      self.registerForTraitChanges([DangerStyleTrait.self]) { (self: Self, _) in
            // Do some styling changes if danger is true or not
      }
    }
}

I want to use these components inside others:

class InfoViewController: UIViewController {
    override func viewDidLoad() {
        super.viewDidLoad()
        
        let button = CustomButton()
        view.addSubview(button)
        // <Button layout elided>
    }
}

In my usual scenario, the custom trait would be applied before presentation:

func showWarningAlert() {
    let viewController = InfoViewController
    viewController.traitOverrides.dangerStyleEnabled = true
    present(viewController, animated: true)
}

But the above does not work! The button never receives the trait change update and so the danger styling never gets applied. I would have assumed that as soon as the view controller's view gets added to the hierarchy, the update would get sent to the button.

If this is expected behaviour, is there a workaround for the above scenario? I would hope it is a reasonable use case to want to set up a whole view controller and then later apply a trait and have it propagate all the way down on presentation.

Thanks!

Answer 1

This isn't specific to custom traits.

The trait isn't changing. It's there from the start. Your implementation of CustomButton should be checking its trait collection right from the start, in addition to checking for changes.

Whatever setup code is in your CustomButton class needs to take into account whatever traits affects its look right from the start.

Registering for a specific trait change is done to allow the view to update itself later if a trait changes.