Utility to create custom VisualElements that correspond to specific types in Unity

ControlUtility.cs

using UnityEngine;
using UnityEditor;
using UnityEngine.UIElements;
using System;
using System.Collections.Generic;

public static class ControlUtility
{
    private static class ControlCache<TBaseControl, TBaseObject> where TBaseControl : VisualElement where TBaseObject : class
    {
        // Key is object type, value is control type.
        public static readonly Dictionary<Type, Type> cache = new Dictionary<Type, Type>();
    }

    public static TBaseControl CreateControl<TBaseControl, TBaseObject>(Type objectType) where TBaseControl : VisualElement where TBaseObject : class
    {
        var baseObjectType = typeof(TBaseObject);

        if (!baseObjectType.IsAssignableFrom(objectType))
        {
            Debug.LogError("Object Type is not compatible with Base Object Type");
            return null;
        }

        var cache = ControlCache<TBaseControl, TBaseObject>.cache;
        if (!cache.TryGetValue(objectType, out Type controlType))
        {
            var controlTypes = TypeCache.GetTypesDerivedFrom<TBaseControl>();

            controlType = FindControlType(controlTypes, objectType, baseObjectType);
            cache[objectType] = controlType;
        }

        if (controlType == null)
            return null;

        return Activator.CreateInstance(controlType) as TBaseControl;
    }

    private static Type FindControlType(IList<Type> controlTypes, Type objectType, Type baseObjectType)
    {
        for (int i = 0; i < controlTypes.Count; i++)
        {
            var ct = controlTypes[i];
            // We force a parameterless constructor to ensure that we can create an instance of the control.
            if (ct.IsAbstract || ct.GetConstructor(Type.EmptyTypes) == null)
                continue;

            var attributes = ct.GetCustomAttributes(typeof(CustomControlAttribute), false);

            foreach (var a in attributes)
                if ((a as CustomControlAttribute).inspectedType == objectType)
                    return ct;                    
        }

        // If we didn't find a control type, we search recursively up the inheritance chain, until reaching the base type.
        if (baseObjectType != objectType)
            return FindControlType(controlTypes, objectType.BaseType, baseObjectType);

        return null;
    }
}

[AttributeUsage(AttributeTargets.Class, Inherited = false, AllowMultiple = true)]
public class CustomControlAttribute : Attribute
{
    public Type inspectedType { get; }
    public CustomControlAttribute(Type inspectedType)
    {
        this.inspectedType = inspectedType;
    }
}

~BasicExample.cs

// Suppose we have a custom base node class and a child node class:
public abstract class NodeBase { }

public class ChildNode : NodeBase { }

//We then can have a base NodeControl to be shown inside some interface for each of its nodes:
[CustomControl(typeof(NodeBase))]
public class NodeControl : VisualElement
{
    public abstract void Init(NodeBase node)
    {
        //Do your thing.
    }
}

// The NodeControl will work as a fallback for all node types, but we also want a specific UI for ChildNode:
[CustomControl(typeof(ChildNode))]
public class ChildNodeControl : NodeControl
{
    public override void Init(NodeBase node)
    {
        //Do your thing.
    }
}

// Then, to create the appropriate NodeControl for a given node we can do this:
var control = ControlUtility.CreateControl<NodeControl, NodeBase>(myNode.GetType());
control.Init(myNode);