From 5d06843fcfe8a7c620ee3b31b82edc930433ad81 Mon Sep 17 00:00:00 2001
From: Guilherme Sousa <guilherme.sousa1994@gmail.com>
Date: Sat, 8 Oct 2022 00:16:54 +0200
Subject: Added custom node export

---
 editor/editor_node.cpp               | 21 +++++++++++++++++++++
 editor/editor_node.h                 |  2 ++
 editor/editor_properties.cpp         |  3 ++-
 editor/scene_tree_editor.cpp         |  5 +++--
 modules/gdscript/gdscript_parser.cpp | 16 +++++++++++-----
 5 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/editor/editor_node.cpp b/editor/editor_node.cpp
index bf50efc4f9..3e50e307f2 100644
--- a/editor/editor_node.cpp
+++ b/editor/editor_node.cpp
@@ -4392,6 +4392,27 @@ Ref<Texture2D> EditorNode::get_class_icon(const String &p_class, const String &p
 	return nullptr;
 }
 
+bool EditorNode::is_object_of_custom_type(const Object *p_object, const StringName &p_class) {
+	ERR_FAIL_COND_V(!p_object, false);
+
+	Ref<Script> scr = p_object->get_script();
+	if (scr.is_null() && Object::cast_to<Script>(p_object)) {
+		scr = p_object;
+	}
+
+	if (scr.is_valid()) {
+		Ref<Script> base_script = scr;
+		while (base_script.is_valid()) {
+			StringName name = EditorNode::get_editor_data().script_class_get_name(base_script->get_path());
+			if (name == p_class) {
+				return true;
+			}
+			base_script = base_script->get_base_script();
+		}
+	}
+	return false;
+}
+
 void EditorNode::progress_add_task(const String &p_task, const String &p_label, int p_steps, bool p_can_cancel) {
 	if (singleton->cmdline_export_mode) {
 		print_line(p_task + ": begin: " + p_label + " steps: " + itos(p_steps));
diff --git a/editor/editor_node.h b/editor/editor_node.h
index 9452425dc8..b5d844558e 100644
--- a/editor/editor_node.h
+++ b/editor/editor_node.h
@@ -827,6 +827,8 @@ public:
 	Ref<Texture2D> get_object_icon(const Object *p_object, const String &p_fallback = "Object");
 	Ref<Texture2D> get_class_icon(const String &p_class, const String &p_fallback = "Object") const;
 
+	bool is_object_of_custom_type(const Object *p_object, const StringName &p_class);
+
 	void show_accept(const String &p_text, const String &p_title);
 	void show_save_accept(const String &p_text, const String &p_title);
 	void show_warning(const String &p_text, const String &p_title = TTR("Warning!"));
diff --git a/editor/editor_properties.cpp b/editor/editor_properties.cpp
index 332e47dc52..85cfd99ab6 100644
--- a/editor/editor_properties.cpp
+++ b/editor/editor_properties.cpp
@@ -3697,7 +3697,8 @@ bool EditorPropertyNodePath::is_drop_valid(const Dictionary &p_drag_data) const
 	}
 
 	for (const StringName &E : valid_types) {
-		if (dropped_node->is_class(E)) {
+		if (dropped_node->is_class(E) ||
+				EditorNode::get_singleton()->is_object_of_custom_type(dropped_node, E)) {
 			return true;
 		}
 	}
diff --git a/editor/scene_tree_editor.cpp b/editor/scene_tree_editor.cpp
index da727a127a..c8bc189106 100644
--- a/editor/scene_tree_editor.cpp
+++ b/editor/scene_tree_editor.cpp
@@ -483,8 +483,9 @@ void SceneTreeEditor::_add_nodes(Node *p_node, TreeItem *p_parent) {
 
 	if (valid_types.size()) {
 		bool valid = false;
-		for (int i = 0; i < valid_types.size(); i++) {
-			if (p_node->is_class(valid_types[i])) {
+		for (const StringName &E : valid_types) {
+			if (p_node->is_class(E) ||
+					EditorNode::get_singleton()->is_object_of_custom_type(p_node, E)) {
 				valid = true;
 				break;
 			}
diff --git a/modules/gdscript/gdscript_parser.cpp b/modules/gdscript/gdscript_parser.cpp
index bdf6fb35b6..4279edf394 100644
--- a/modules/gdscript/gdscript_parser.cpp
+++ b/modules/gdscript/gdscript_parser.cpp
@@ -3765,13 +3765,19 @@ bool GDScriptParser::export_annotations(const AnnotationNode *p_annotation, Node
 				break;
 			case GDScriptParser::DataType::CLASS:
 				// Can assume type is a global GDScript class.
-				if (!ClassDB::is_parent_class(export_type.native_type, SNAME("Resource"))) {
-					push_error(R"(Exported script type must extend Resource.)");
+				if (ClassDB::is_parent_class(export_type.native_type, SNAME("Resource"))) {
+					variable->export_info.type = Variant::OBJECT;
+					variable->export_info.hint = PROPERTY_HINT_RESOURCE_TYPE;
+					variable->export_info.hint_string = export_type.class_type->identifier->name;
+				} else if (ClassDB::is_parent_class(export_type.native_type, SNAME("Node"))) {
+					variable->export_info.type = Variant::OBJECT;
+					variable->export_info.hint = PROPERTY_HINT_NODE_TYPE;
+					variable->export_info.hint_string = export_type.class_type->identifier->name;
+				} else {
+					push_error(R"(Export type can only be built-in, a resource, a node or an enum.)", variable);
 					return false;
 				}
-				variable->export_info.type = Variant::OBJECT;
-				variable->export_info.hint = PROPERTY_HINT_RESOURCE_TYPE;
-				variable->export_info.hint_string = export_type.class_type->identifier->name;
+
 				break;
 			case GDScriptParser::DataType::SCRIPT: {
 				StringName class_name;
-- 
cgit v1.2.3