How do I determine the name and reference of a class being targeted by an attribute at runtime?

Question

I am trying to create a custom class attribute for Unity that prevents a targeted MonoBehaviour from existing on more than one object in a scene. I have done some searching and it's been said that in order to get the Type of class the attribute is targeting, I should use the constructor of the attribute; it can't be done using reflection

[AttributeUsage(AttributeTargets.Class, AllowMultiple = false, Inherited = true)]
public class DisallowMultipleInScene : Attribute
{
    public DisallowMultipleInScene(Type type)
    {
        // Do stuff with type here
    }
}

If I'm not mistaken, this means that using it on a class, such as one named ManagerManager, would be achieved like this:

[DisallowMultipleInScene(typeof(ManagerManager))]
public class ManagerManager : MonoBehaviour
{
    // Implementation
}

This seems a bit redundant, and would also allow passing in the incorrect class name. For disallowing multiple components (classes inheriting from MonoBehaviour) being placed on the same object, the [DisallowMultipleComponent] attribute is used. This attribute is similar to what I'd like. You are not required to pass in the name of the class it is being applied to, it just seems to know.

I had a look at the source in the UnityCsReference GitHub, in an attempt to learn how it works behind the scenes, but there appeared to be no implementation, just a definition, located here, excerpted below:

[RequiredByNativeCode]
[AttributeUsage(AttributeTargets.Class, Inherited = false)]
public sealed class DisallowMultipleComponent : Attribute {}

So for my own implementation, to work similarly to the [DisallowMultipleComponent] attribute, my attribute would need to determine the class it is being applied to, and get a reference to the MonoBehaviour script, so that it can be removed from the object it was just added to.

So, firstly, how does the [DisallowMultipleComponent] attribute get around the requirement of passing the class type in as an attribute parameter, and how can I do it?

And, secondly, how do I get a reference to the newly-created instance of the class annotated with the attribute?

Answer 1

There are four steps to achieving the functionality you want.

1. Detect all types with the appropriate attribute. You can do this by looping through each Assembly in AppDomain.CurrentDomain. You'll want to cache these types each time your script assembly reloads, which you can check with a static class and the InitializeOnLoad attribute from the editor. (You definitely don't want to be doing reflection if you don't have to).
2. Detect when objects are added/modified in the scene hierarchy. This can be accomplished with the event EditorApplication.hierarchyChanged.
3. Check if any component has been added to the scene that shouldn't be there. This can be accomplished with the UnityEditor.SceneManagement.EditorSceneManager class by looping over all root objects in the scene and tracking the appropriate information.
4. Decide what to do if you encounter multiple of the same component (destroy, show a message to the user, etc.). This one is sort of up to you, but I've included a logical answer below.

This can be achieved with the following attribute

using System;

[AttributeUsage(AttributeTargets.Class, AllowMultiple = false, Inherited = true)]
public class DisallowMultipleComponentsInSceneAttribute : Attribute
{
    public DisallowMultipleComponentsInSceneAttribute()
    {

    }
}

and the following Editor script, which must be placed inside of an "Editor" folder within your project.

using UnityEngine;
using UnityEngine.SceneManagement;

using UnityEditor;
using UnityEditor.SceneManagement;

using System;
using System.Reflection;
using System.Collections.Generic;

[InitializeOnLoad]
public static class SceneHierarchyMonitor 
{
    private class TrackingData
    {
        public Dictionary<Scene, Component> components = new Dictionary<Scene, Component>();
    }

    private static Dictionary<Type, TrackingData> __trackingData = new Dictionary<Type, TrackingData>();

    static SceneHierarchyMonitor()
    {
        EditorApplication.hierarchyChanged += OnHierarchyChanged;

        foreach (Assembly assembly in AppDomain.CurrentDomain.GetAssemblies())
        {
            foreach (Type type in assembly.GetTypes())
            {
                if (type.GetCustomAttribute<DisallowMultipleComponentsInSceneAttribute>() != null)
                {
                    __trackingData.Add(type, new TrackingData());
                }
            }
        }

        for (int i = 0; i < EditorSceneManager.sceneCount; ++i)
        {
            MonitorScene(EditorSceneManager.GetSceneAt(i));
        }
    }

    private static void OnHierarchyChanged()
    {
        for (int i = 0; i < EditorSceneManager.sceneCount; ++i)
        {
            MonitorScene(EditorSceneManager.GetSceneAt(i));
        }
    }

    private static void MonitorScene(Scene scene)
    {
        foreach (KeyValuePair<Type, TrackingData> kvp in __trackingData)
        {
            // If the scene hasn't been tracked, initialize the component to a null value.
            bool isOpeningScene = false;
            if (!kvp.Value.components.ContainsKey(scene))
            {
                isOpeningScene = true;
                kvp.Value.components[scene] = null;
            }

            foreach (GameObject rootGameObject in scene.GetRootGameObjects())
            {
                Component[] components = rootGameObject.GetComponentsInChildren(kvp.Key, true);
                for (int i = 0; i < components.Length; ++i)
                {
                    Component component = components[i];

                    // If we haven't found a component of this type yet, set it to remember it. This will occur when either:
                    // 1. The component is added for the first time in a given scene.
                    // 2. The scene is being opened and we didn't have any tracking data previously.
                    if (kvp.Value.components[scene] == null)
                    {
                        kvp.Value.components[scene] = component;
                    }
                    else
                    {
                        // You can determine what to do with extra components. This makes sense to me, but you can change the
                        // behavior as you see fit.
                        if (kvp.Value.components[scene] != component)
                        {
                            GameObject gameObject = component.gameObject;
                            EditorGUIUtility.PingObject(gameObject);
                            if (!isOpeningScene)
                            {
                                Debug.LogError($"Destroying \"{component}\" because it has the attribute \"{typeof(DisallowMultipleComponentsInSceneAttribute).Name}\", " +
                                    $"and one of these components already exists in scene \"{scene.name}.\"", gameObject);
                                GameObject.DestroyImmediate(component);
                                EditorUtility.SetDirty(gameObject);
                            }
                            else
                            {
                                Debug.LogWarning($"Found multiple components of type {kvp.Key.Name} in scene {scene.name}. Please ensure there is exactly one " +
                                    $"instance of this type in the scene before continuing.", component.gameObject);
                            }
                        }
                    }
                }
            }
        }
    }
}

I've tested this out in the Editor and it works quite nicely.

A final note:

• This Editor script works fine on a small project, but on a large project with big scenes and/or lots of script files, it's possible to run into performance issues. Even a few fractions of a second delay when you reload your script assembly is noticeable, and it can add up over the lifetime of the project.