From a3309215c200dff1cb8930cbdf83dbc0cff6d491 Mon Sep 17 00:00:00 2001
From: kobewi <kobewi4e@gmail.com>
Date: Mon, 28 Feb 2022 01:57:34 +0100
Subject: [PATCH] Improve handling of custom types

---
 doc/classes/EditorPlugin.xml    |  1 +
 editor/create_dialog.cpp        |  4 ++--
 editor/editor_data.cpp          | 26 ++++++++++++++++++++++++++
 editor/editor_data.h            |  3 +++
 editor/editor_inspector.cpp     | 10 ++++++++--
 editor/editor_node.cpp          | 12 ++----------
 editor/script_create_dialog.cpp | 12 ++++++++++--
 7 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/doc/classes/EditorPlugin.xml b/doc/classes/EditorPlugin.xml
index fdd3807b698b..9271e384872b 100644
--- a/doc/classes/EditorPlugin.xml
+++ b/doc/classes/EditorPlugin.xml
@@ -406,6 +406,7 @@
 				When a given node or resource is selected, the base type will be instantiated (e.g. "Node3D", "Control", "Resource"), then the script will be loaded and set to this object.
 				You can use the virtual method [method _handles] to check if your custom object is being edited by checking the script or using the [code]is[/code] keyword.
 				During run-time, this will be a simple object with a script so this function does not need to be called then.
+				[b]Note:[/b] Custom types added this way are not true classes. They are just a helper to create a node with specific script.
 			</description>
 		</method>
 		<method name="add_debugger_plugin">
diff --git a/editor/create_dialog.cpp b/editor/create_dialog.cpp
index 8ccfda1145cf..2d70934c9c83 100644
--- a/editor/create_dialog.cpp
+++ b/editor/create_dialog.cpp
@@ -665,7 +665,7 @@ void CreateDialog::_save_and_update_favorite_list() {
 			for (int i = 0; i < favorite_list.size(); i++) {
 				String l = favorite_list[i];
 				String name = l.get_slicec(' ', 0);
-				if (!(ClassDB::class_exists(name) || ScriptServer::is_global_class(name))) {
+				if (!EditorNode::get_editor_data().is_type_recognized(name)) {
 					continue;
 				}
 				f->store_line(l);
@@ -692,7 +692,7 @@ void CreateDialog::_load_favorites_and_history() {
 			String l = f->get_line().strip_edges();
 			String name = l.get_slicec(' ', 0);
 
-			if ((ClassDB::class_exists(name) || ScriptServer::is_global_class(name)) && !_is_class_disabled_by_feature_profile(name)) {
+			if (EditorNode::get_editor_data().is_type_recognized(name) && !_is_class_disabled_by_feature_profile(name)) {
 				recent->add_item(l, EditorNode::get_singleton()->get_class_icon(name, icon_fallback));
 			}
 		}
diff --git a/editor/editor_data.cpp b/editor/editor_data.cpp
index 231ae198d207..227e1b28e2b2 100644
--- a/editor/editor_data.cpp
+++ b/editor/editor_data.cpp
@@ -509,6 +509,32 @@ Variant EditorData::instance_custom_type(const String &p_type, const String &p_i
 	return Variant();
 }
 
+const EditorData::CustomType *EditorData::get_custom_type_by_name(const String &p_type) const {
+	for (const KeyValue<String, Vector<CustomType>> &E : custom_types) {
+		for (const CustomType &F : E.value) {
+			if (F.name == p_type) {
+				return &F;
+			}
+		}
+	}
+	return nullptr;
+}
+
+const EditorData::CustomType *EditorData::get_custom_type_by_path(const String &p_path) const {
+	for (const KeyValue<String, Vector<CustomType>> &E : custom_types) {
+		for (const CustomType &F : E.value) {
+			if (F.script->get_path() == p_path) {
+				return &F;
+			}
+		}
+	}
+	return nullptr;
+}
+
+bool EditorData::is_type_recognized(const String &p_type) const {
+	return ClassDB::class_exists(p_type) || ScriptServer::is_global_class(p_type) || get_custom_type_by_name(p_type);
+}
+
 void EditorData::remove_custom_type(const String &p_type) {
 	for (KeyValue<String, Vector<CustomType>> &E : custom_types) {
 		for (int i = 0; i < E.value.size(); i++) {
diff --git a/editor/editor_data.h b/editor/editor_data.h
index 1da188c546c2..4f1740d4f073 100644
--- a/editor/editor_data.h
+++ b/editor/editor_data.h
@@ -184,6 +184,9 @@ class EditorData {
 	Variant instance_custom_type(const String &p_type, const String &p_inherits);
 	void remove_custom_type(const String &p_type);
 	const HashMap<String, Vector<CustomType>> &get_custom_types() const { return custom_types; }
+	const CustomType *get_custom_type_by_name(const String &p_name) const;
+	const CustomType *get_custom_type_by_path(const String &p_path) const;
+	bool is_type_recognized(const String &p_type) const;
 
 	void instantiate_object_properties(Object *p_object);
 
diff --git a/editor/editor_inspector.cpp b/editor/editor_inspector.cpp
index bdf1c314f862..5bd08d1e633a 100644
--- a/editor/editor_inspector.cpp
+++ b/editor/editor_inspector.cpp
@@ -2738,7 +2738,7 @@ void EditorInspector::update_tree() {
 			doc_name = p.name;
 
 			// Set the category icon.
-			if (!ClassDB::class_exists(type) && !ScriptServer::is_global_class(type) && p.hint_string.length() && FileAccess::exists(p.hint_string)) {
+			if (!EditorNode::get_editor_data().is_type_recognized(type) && p.hint_string.length() && FileAccess::exists(p.hint_string)) {
 				// If we have a category inside a script, search for the first script with a valid icon.
 				Ref<Script> script = ResourceLoader::load(p.hint_string, "Script");
 				StringName base_type;
@@ -2757,10 +2757,16 @@ void EditorInspector::update_tree() {
 				while (script.is_valid()) {
 					name = EditorNode::get_editor_data().script_class_get_name(script->get_path());
 					String icon_path = EditorNode::get_editor_data().script_class_get_icon_path(name);
-					if (name != StringName() && icon_path.length()) {
+					if (name != StringName() && !icon_path.is_empty()) {
 						category->icon = ResourceLoader::load(icon_path, "Texture");
 						break;
 					}
+
+					const EditorData::CustomType *ctype = EditorNode::get_editor_data().get_custom_type_by_path(script->get_path());
+					if (ctype) {
+						category->icon = ctype->icon;
+						break;
+					}
 					script = script->get_base_script();
 				}
 				if (category->icon.is_null() && has_theme_icon(base_type, SNAME("EditorIcons"))) {
diff --git a/editor/editor_node.cpp b/editor/editor_node.cpp
index 58377eef0f1a..c328420a5f7d 100644
--- a/editor/editor_node.cpp
+++ b/editor/editor_node.cpp
@@ -4338,16 +4338,8 @@ Ref<Texture2D> EditorNode::get_class_icon(const String &p_class, const String &p
 		}
 	}
 
-	const HashMap<String, Vector<EditorData::CustomType>> &p_map = EditorNode::get_editor_data().get_custom_types();
-	for (const KeyValue<String, Vector<EditorData::CustomType>> &E : p_map) {
-		const Vector<EditorData::CustomType> &ct = E.value;
-		for (int i = 0; i < ct.size(); ++i) {
-			if (ct[i].name == p_class) {
-				if (ct[i].icon.is_valid()) {
-					return ct[i].icon;
-				}
-			}
-		}
+	if (const EditorData::CustomType *ctype = EditorNode::get_editor_data().get_custom_type_by_name(p_class)) {
+		return ctype->icon;
 	}
 
 	if (gui_base->has_theme_icon(p_class, SNAME("EditorIcons"))) {
diff --git a/editor/script_create_dialog.cpp b/editor/script_create_dialog.cpp
index 7d065b4920d9..26e79135e007 100644
--- a/editor/script_create_dialog.cpp
+++ b/editor/script_create_dialog.cpp
@@ -202,7 +202,7 @@ bool ScriptCreateDialog::_validate_parent(const String &p_string) {
 		}
 	}
 
-	return ClassDB::class_exists(p_string) || ScriptServer::is_global_class(p_string);
+	return EditorNode::get_editor_data().is_type_recognized(p_string);
 }
 
 bool ScriptCreateDialog::_validate_class(const String &p_string) {
@@ -372,7 +372,15 @@ void ScriptCreateDialog::_create_new() {
 
 	const ScriptLanguage::ScriptTemplate sinfo = _get_current_template();
 
-	scr = ScriptServer::get_language(language_menu->get_selected())->make_template(sinfo.content, cname_param, parent_name->get_text());
+	String parent_class = parent_name->get_text();
+	if (!ClassDB::class_exists(parent_class) && !ScriptServer::is_global_class(parent_class)) {
+		// If base is a custom type, replace with script path instead.
+		const EditorData::CustomType *type = EditorNode::get_editor_data().get_custom_type_by_name(parent_class);
+		ERR_FAIL_NULL(type);
+		parent_class = "\"" + type->script->get_path() + "\"";
+	}
+
+	scr = ScriptServer::get_language(language_menu->get_selected())->make_template(sinfo.content, cname_param, parent_class);
 
 	if (has_named_classes) {
 		String cname = class_name->get_text();