[UnityNotNull] attribute

UnityNotNull.cs

using System;
using System.Linq;
using System.Reflection;
using UnityEngine;
#if UNITY_EDITOR
using UnityEditor;
using UnityEditor.Callbacks;
#endif

namespace com.partlyhuman
{
    /// <summary>
    /// For object references that need to be assigned in Unity Editor's inspector.
    /// Annotate fields with this flag if they are required to be assigned.
    /// The inspector will color these fields red if not assigned, and will throw an error if you
    /// enter playmode with any of these references null. 
    /// Click the error to identify the source object.
    /// </summary>
    public class UnityNotNull : PropertyAttribute
    {
        // Should only be used on object-reference fields.
        // Ideally use a unique component type to identify the proper reference.
        public static bool IsAttributeValidFor(FieldInfo field)
        {
            return field.FieldType.IsSubclassOf(typeof(UnityEngine.Object));
        }
    }

    public static partial class MonoBehaviourExtensions
    {
        /// <summary>
        /// In the editor, assert all [UnityNotNull] fields in this instance are defined.
        /// In builds, does nothing.
        /// You can run this on demand but typically this is run at play time in the editor by EditorUnityNotNullEnforcement
        /// </summary>
        public static void AssertUnityNotNullFields(this MonoBehaviour source)
        {
#if UNITY_EDITOR
            var fields = source.GetType()
                .GetFields(BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance)
                .Where(f => Attribute.GetCustomAttribute(f, typeof(UnityNotNull)) != null);
            foreach (FieldInfo field in fields)
            {
                if (!UnityNotNull.IsAttributeValidFor(field))
                {
                    Debug.LogWarningFormat(source, "[UnityNotNull] not allowed for fields of type {0} in {1}::{2}", 
                        field.FieldType.Name, field.DeclaringType.Name, field.Name);
                    continue;
                }
                // Debug.LogFormat("Checking {0}", field.Name);
                Debug.AssertFormat(field.GetValue(source) != null, source, 
                    "[UnityNotNull] attributed field is null {0}::{1}", field.DeclaringType.Name, field.Name);
            }
#endif
        }
    }

#if UNITY_EDITOR
    public static class EditorUnityNotNullEnforcement
    {
        // Runs when you hit play
        [DidReloadScripts]
        public static void OnScriptsReloaded()
        {
            EnforceAllUnityNotNullFields();
        }

        [MenuItem("GameObject/Check for Null References")]
        public static void EnforceAllUnityNotNullFields()
        {
            foreach (var go in GameObject.FindObjectsOfType<MonoBehaviour>())
            {
                go.AssertUnityNotNullFields();
            }
        }
    }
#endif

#if UNITY_EDITOR
    [CustomPropertyDrawer(typeof(UnityNotNull))]
    public class AssertNotNullInspector : PropertyDrawer
    {
        static readonly Color UNASSIGNED_BG = Color.red;
        const float BUTTON_W = 46f;
        const string FIND_IN_SCENE = "Scene";
        const string FIND_IN_ASSETS = "Prefab";

        public override void OnGUI(Rect position, SerializedProperty property, GUIContent label)
        {
            if (!UnityNotNull.IsAttributeValidFor(fieldInfo) || property.objectReferenceValue != null)
            {
                EditorGUI.PropertyField(position, property, label);
            }
            else
            {
                Color oldColor = GUI.backgroundColor;
                GUI.backgroundColor = UNASSIGNED_BG;

                position.width -= 2 * BUTTON_W;
                var buttonPos = new Rect
                {
                    x = position.x + position.width,
                    y = position.y,
                    width = BUTTON_W,
                    height = position.height,
                };

                EditorGUI.PropertyField(position, property, label);
                GUI.backgroundColor = oldColor;

                if (GUI.Button(buttonPos, FIND_IN_SCENE))
                {
                    property.objectReferenceValue = GameObject.FindObjectOfType(fieldInfo.FieldType);
                }

                buttonPos.x += buttonPos.width;
                if (GUI.Button(buttonPos, FIND_IN_ASSETS))
                {
                    property.objectReferenceValue = AssetDatabase.FindAssets("t:Prefab")
                        .Select(guid => AssetDatabase.LoadAssetAtPath<GameObject>(AssetDatabase.GUIDToAssetPath(guid)))
                        .FirstOrDefault(asset => asset && asset.GetComponent(fieldInfo.FieldType) != null);
                    Resources.UnloadUnusedAssets();
                }
            }
        }
    }
#endif
}