1 files changed, 45 insertions, 21 deletions

diff --git a/addons/ClassExporter/Plugin.cs b/addons/ClassExporter/Plugin.cs
index 8a7e349..e945f63 100644
--- a/addons/ClassExporter/Plugin.cs
+++ b/addons/ClassExporter/Plugin.cs
@@ -13,7 +13,7 @@ namespace Texty.addons.ClassExporter
     {
         private const string LoadTypesMenuItem = "Load types";
         private const string ReloadTypesMenuItem = "Reload types";
-        private readonly List<Type> _types = new List<Type>();
+        private readonly List<TypeEntry> _registeredTypes = new List<TypeEntry>();
 
         public override void _EnterTree()
         {
@@ -23,6 +23,16 @@ namespace Texty.addons.ClassExporter
             RegisterTypes();
         }
 
+        public override void EnablePlugin()
+        {
+            RegisterTypes();
+        }
+
+        public override void DisablePlugin()
+        {
+            DeRegisterTypes();
+        }
+        
         public override void _ExitTree()
         {
             DeRegisterTypes();
@@ -31,6 +41,7 @@ namespace Texty.addons.ClassExporter
             Disconnect("resource_saved", this, nameof(OnResourceSaved));
         }
 
+        // Required by Godot
         // ReSharper disable once UnusedParameter.Local
         private void OnResourceSaved(Resource resource)
         {
@@ -49,45 +60,58 @@ namespace Texty.addons.ClassExporter
             var newTypes = from type in assembly.GetTypes()
                 where type.IsSubclassOf(typeof(Resource)) || type.IsSubclassOf(typeof(Node))
                 where !type.IsSubclassOf(typeof(EditorPlugin))
-                where !type.IsAbstract
-                where !_types.Contains(type)
-                select new RegistrableType(type);
+                where _registeredTypes.Find(r => r.Equals(type)) == null
+                select new TypeEntry(type);
 
             foreach (var type in newTypes)
             {
-                GD.Print($"Exporting class {type.Name}: {type.Script.ResourcePath}");
+                GD.Print($"Exporting class {type.Name}");
                 AddCustomType(type.Name, type.Base, type.Script, null);
-                _types.Add(type.Type);
+                _registeredTypes.Add(type);
             }
         }
 
         private void DeRegisterTypes()
         {
-            _types.ForEach(type => RemoveCustomType(type.Name));
-            _types.Clear();
+            foreach (var type in _registeredTypes)
+            {
+                GD.Print($"Removing class {type.Name}");
+                RemoveCustomType(type.Name);
+            }
+
+            _registeredTypes.Clear();
         }
 
-        private readonly struct RegistrableType
+        private class TypeEntry : Reference, IEquatable<Type>
         {
-            public RegistrableType(Type type)
+            public readonly string Base = "";
+            public readonly string Name = "";
+            public readonly Script Script = new CSharpScript();
+
+            // Required by Godot
+            // ReSharper disable once UnusedMember.Local
+            public TypeEntry()
             {
-                Type = type;
             }
 
-            public readonly Type Type;
+            public TypeEntry(Type type)
+            {
+                Name = type.Name;
+                Base = type.BaseType?.Name ?? "Object";
+                Script = LoadScript(type.FullName ?? "");
+            }
 
-            public CSharpScript Script
+            public bool Equals(Type other)
             {
-                get
-                {
-                    var components = (Type.FullName ?? "").Split('.');
-                    var path = string.Join("/", components.Skip(1)) + ".cs";
-                    return ResourceLoader.Load<CSharpScript>(path);
-                }
+                return other != null && Name == other.Name && Base == other.BaseType?.Name;
             }
 
-            public string Name => Type.Name;
-            public string Base => Type.BaseType?.Name ?? "";
+            private static CSharpScript LoadScript(string fullName)
+            {
+                var components = fullName.Split('.');
+                var path = string.Join("/", components.Skip(1)) + ".cs";
+                return ResourceLoader.Load<CSharpScript>(path);
+            }
         }
     }
 }