Workaround for Unity's inability to keep an input field selected when you click a button

InputFieldActivationRestorer.cs

using System.Reflection;
using TMPro;
using UnityEngine;
using UnityEngine.EventSystems;
using UnityEngine.UI;

// Workaround for Unity's inability to keep an input field selected when you click a button.
// Call ReactivateInputField to immediately activate and restore caret position/selection.
// This class also tracks the state of the last selected input field in Update, allowing easier restoration of the input field when it becomes deselected.
// Common use cases are to:
// - restore it immediately when it becomes deselected (in the callback for inputField.onEndEdit)
// - restore it immediately once the user clicks a particular button
// - restore it once the user returns to a particular state in your app
public class InputFieldActivationRestorer : MonoSingleton<InputFieldActivationRestorer> {
    // The current restorable settings 
    public RestorableInputFieldSettings lastSelectedInputFieldSettings;
    public bool lastSelectedInputFieldIsCurrentlySelected => lastSelectedInputFieldSettings.isRestorable && EventSystem.current != null && EventSystem.current.currentSelectedGameObject == lastSelectedInputFieldSettings.restorableGameObject;
    
    
    // This struct contains all the properties that can be restored
    [System.Serializable]
    public struct RestorableInputFieldSettings {
        public bool isRestorable => restorableInputField != null || restorableTMPInputField != null;
        public GameObject restorableGameObject {
            get {
                if (restorableInputField != null) return restorableInputField.gameObject;
                else if (restorableTMPInputField != null) return restorableTMPInputField.gameObject;
                else return null;
            }
        }
        
        public InputField restorableInputField;
        public TMP_InputField restorableTMPInputField;
        public int caretPosition;
        public int selectionAnchorPosition;
        public int selectionFocusPosition;
        
        
        // Sets the input field to be restored automatically using whatever input field is currently selected
        public static RestorableInputFieldSettings FromCurrentSelection() {
            var selectedInputField = GetSelectedUGUIInputField();
            if (selectedInputField != null) {
                return new RestorableInputFieldSettings(selectedInputField);
            }
            var selectedTMPInputField = GetSelectedTMPInputField();
            if(selectedTMPInputField != null) {
                return new RestorableInputFieldSettings(selectedTMPInputField);
            }

            return default;

            static TMP_InputField GetSelectedTMPInputField () {
                if(EventSystem.current == null) return null;
                var selectedGO = EventSystem.current.currentSelectedGameObject;
                if (selectedGO == null) return null;
                var inputField = selectedGO.GetComponentInParent<TMP_InputField>();
                return inputField != null && inputField.isFocused ? inputField : null;
            }
    
            static InputField GetSelectedUGUIInputField () {
                if(EventSystem.current == null) return null;
                var selectedGO = EventSystem.current.currentSelectedGameObject;
                if (selectedGO == null) return null;
                var inputField = selectedGO.GetComponentInParent<InputField>();
                return inputField != null && inputField.isFocused ? inputField : null;
            }
        }

        public RestorableInputFieldSettings(InputField inputField) {
            restorableInputField = inputField;
            restorableTMPInputField = null;
            caretPosition = inputField.caretPosition;
            selectionAnchorPosition = inputField.selectionAnchorPosition;
            selectionFocusPosition = inputField.selectionFocusPosition;
        }
    
        public RestorableInputFieldSettings(TMP_InputField tmpInputField) {
            restorableInputField = null;
            restorableTMPInputField = tmpInputField;
            caretPosition = tmpInputField.caretPosition;
            selectionAnchorPosition = tmpInputField.selectionAnchorPosition;
            selectionFocusPosition = tmpInputField.selectionFocusPosition;
        }
        
        public override bool Equals(object obj) {
            if (obj == null || GetType() != obj.GetType()) {
                return false;
            }

            RestorableInputFieldSettings other = (RestorableInputFieldSettings)obj;
            return (restorableInputField == other.restorableInputField &&
                    restorableTMPInputField == other.restorableTMPInputField &&
                    caretPosition == other.caretPosition &&
                    selectionAnchorPosition == other.selectionAnchorPosition &&
                    selectionFocusPosition == other.selectionFocusPosition);
        }

        public override int GetHashCode() {
            unchecked {
                int hash = 17;
                hash = hash * 23 + (restorableInputField != null ? restorableInputField.GetHashCode() : 0);
                hash = hash * 23 + (restorableTMPInputField != null ? restorableTMPInputField.GetHashCode() : 0);
                hash = hash * 23 + caretPosition.GetHashCode();
                hash = hash * 23 + selectionAnchorPosition.GetHashCode();
                hash = hash * 23 + selectionFocusPosition.GetHashCode();
                return hash;
            }
        }
    }


    public void ReactivateLastSelectedInputField() {
        ReactivateInputField(lastSelectedInputFieldSettings);
    }

    // Selects the previously selected input field and restores the selection state.
    public static void ReactivateInputField(RestorableInputFieldSettings restorableState) {
        if (!restorableState.isRestorable) return;

        // Copy the struct just in case something here causes it to be updated while we're restoring it
        if (restorableState.restorableInputField != null) {
            restorableState.restorableInputField.ActivateInputField();
            
            // This forces TMP to immediately set itself selectable again.
            MethodInfo mInfoMethod = typeof(InputField).GetMethod("LateUpdate", BindingFlags.Instance | BindingFlags.NonPublic);
            mInfoMethod.Invoke(restorableState.restorableInputField, null);
        }
        else if (restorableState.restorableTMPInputField != null) {
            restorableState.restorableTMPInputField.ActivateInputField();
            
            // This forces TMP to immediately set itself selectable again.
            MethodInfo mInfoMethod = typeof(TMP_InputField).GetMethod("LateUpdate", BindingFlags.Instance | BindingFlags.NonPublic);
            mInfoMethod.Invoke(restorableState.restorableTMPInputField, null);
        }
        
        if (restorableState.restorableInputField != null && EventSystem.current.currentSelectedGameObject == restorableState.restorableInputField.gameObject) {
            restorableState.restorableInputField.caretPosition = restorableState.caretPosition;
            restorableState.restorableInputField.selectionAnchorPosition = restorableState.selectionAnchorPosition;
            restorableState.restorableInputField.selectionFocusPosition = restorableState.selectionFocusPosition;
            restorableState.restorableInputField.ForceLabelUpdate();
        }
        else if (restorableState.restorableTMPInputField != null && EventSystem.current.currentSelectedGameObject == restorableState.restorableTMPInputField.gameObject) {
            restorableState.restorableTMPInputField.caretPosition = restorableState.caretPosition;
            restorableState.restorableTMPInputField.selectionAnchorPosition = restorableState.selectionAnchorPosition;
            restorableState.restorableTMPInputField.selectionFocusPosition = restorableState.selectionFocusPosition;
            restorableState.restorableTMPInputField.ForceLabelUpdate();
        }
    }

    void Update() {
        if (EventSystem.current != null) {
            var currentSelectionState = RestorableInputFieldSettings.FromCurrentSelection();
            if (currentSelectionState.isRestorable) lastSelectedInputFieldSettings = currentSelectionState;
        }
    }
}