using System; using UnityEngine; using UnityEditor; using UnityEditor.UIElements; using UnityEngine.UIElements; using UnityEditor.SceneManagement; using EditorAttributes.Editor.Utility; using Object = UnityEngine.Object; namespace EditorAttributes.Editor { [CustomPropertyDrawer(typeof(RequiredAttribute))] public class RequiredDrawer : PropertyDrawerBase { public override VisualElement CreatePropertyGUI(SerializedProperty property) { var root = new VisualElement(); var requiredAttribute = attribute as RequiredAttribute; var propertyField = CreatePropertyField(property); root.Add(propertyField); if (property.propertyType == SerializedPropertyType.ObjectReference) { var helpBox = new HelpBox($"The field {property.displayName} must be assigned", HelpBoxMessageType.Error); var fixButton = new Button(() => FixNullReference(property, requiredAttribute)) { text = "Fix" }; if (CanApplyGlobalColor) { helpBox.style.color = EditorExtension.GLOBAL_COLOR; helpBox.style.backgroundColor = EditorExtension.GLOBAL_COLOR / 2f; } if (requiredAttribute.FixMode != ReferenceFixMode.None) helpBox.Add(fixButton); DoRequiredCheck(); propertyField.RegisterCallback((changeEvent) => DoRequiredCheck()); void DoRequiredCheck() { if (property.objectReferenceValue == null) { AddElement(root, helpBox); } else { RemoveElement(root, helpBox); } } } else { root.Add(new HelpBox("The attached field must derive from UnityEngine.Object", HelpBoxMessageType.Error)); } return root; } private void FixNullReference(SerializedProperty property, RequiredAttribute requiredAttribute) { var memberType = ReflectionUtility.GetMemberInfoType(ReflectionUtility.GetValidMemberInfo(property.name, property)); foreach (var target in property.serializedObject.targetObjects) { if (!IsReferenceFixValid(property, target.GetType(), requiredAttribute)) return; var component = target as Component; Object objectReference = null; switch (requiredAttribute.FixMode) { case ReferenceFixMode.Auto: for (int i = 0; i <= 3; i++) { if (objectReference != null) break; switch (i) { case 0: GetSelf(); break; case 1: GetChild(); break; case 2: GetParent(); break; case 3: GetScene(); break; } } break; case ReferenceFixMode.Self: GetSelf(); break; case ReferenceFixMode.Children: GetChild(); break; case ReferenceFixMode.Parents: GetParent(); break; case ReferenceFixMode.Scene: GetScene(); break; case ReferenceFixMode.Custom: objectReference = ReflectionUtility.FindFunction(requiredAttribute.CustomFixFunctionName, target).Invoke(target, null) as Object; break; } property.objectReferenceValue = objectReference; property.serializedObject.ApplyModifiedProperties(); if (property.objectReferenceValue == null) Debug.LogWarning($"Could not find a valid reference of the type {memberType} with ReferenceFixMode.{requiredAttribute.FixMode}", target); void GetSelf() { if (typeof(Component).IsAssignableFrom(memberType)) { objectReference = component.GetComponent(memberType); } else if (typeof(GameObject).IsAssignableFrom(memberType)) { objectReference = component.gameObject; } } void GetChild() { if (typeof(Component).IsAssignableFrom(memberType)) { objectReference = component.GetComponentInChildren(memberType, true); } else if (typeof(GameObject).IsAssignableFrom(memberType)) { if (component.transform.childCount > 0) objectReference = component.transform.GetChild(0).gameObject; } } void GetParent() { if (typeof(Component).IsAssignableFrom(memberType)) { objectReference = component.GetComponentInParent(memberType, true); } else if (typeof(GameObject).IsAssignableFrom(memberType)) { if (component.transform.parent != null) objectReference = component.transform.parent.gameObject; } } void GetScene() { if (PrefabStageUtility.GetCurrentPrefabStage() == null) objectReference = Object.FindFirstObjectByType(memberType, FindObjectsInactive.Include); } } } private bool IsReferenceFixValid(SerializedProperty property, Type targetType, RequiredAttribute requiredAttribute) { bool isComponent = typeof(Component).IsAssignableFrom(targetType); bool isScriptableObject = typeof(ScriptableObject).IsAssignableFrom(targetType); var targetObject = property.serializedObject.targetObject; switch (requiredAttribute.FixMode) { case ReferenceFixMode.Auto: case ReferenceFixMode.Self: case ReferenceFixMode.Children: case ReferenceFixMode.Parents: case ReferenceFixMode.Scene: if (isScriptableObject) { Debug.LogError($"{requiredAttribute.FixMode} is not valid on ScriptableObjects", targetObject); return false; } else if (!isComponent) { Debug.LogError($"{requiredAttribute.FixMode} is not valid on the type {targetType}", targetObject); return false; } break; case ReferenceFixMode.Custom: string functionName = requiredAttribute.CustomFixFunctionName; if (string.IsNullOrEmpty(functionName)) { Debug.LogError($"No custom fix function name was provided", targetObject); return false; } var functionInfo = ReflectionUtility.FindFunction(functionName, property); if (functionInfo == null) { Debug.LogError($"Could not find function {functionName}. If this function is inherited make sure is marked at protected", targetObject); return false; } else if (functionInfo.GetParameters().Length != 0) { Debug.LogError($"The function {functionName} cannot have parameters", targetObject); return false; } else if (!typeof(Object).IsAssignableFrom(functionInfo.ReturnType)) { Debug.LogError($"The function {functionName} needs to return a Unity.Object", targetObject); return false; } break; } return true; } } }