1 files changed, 14 insertions, 10 deletions
diff --git a/editor/editor_node.cpp b/editor/editor_node.cpp
index 3a9715d1ab..c6613cdf63 100644
--- a/editor/editor_node.cpp
+++ b/editor/editor_node.cpp
@@ -30,21 +30,21 @@
 
 #include "editor_node.h"
 
-#include "core/bind/core_bind.h"
-#include "core/class_db.h"
+#include "core/config/project_settings.h"
+#include "core/core_bind.h"
 #include "core/input/input.h"
 #include "core/io/config_file.h"
 #include "core/io/image_loader.h"
 #include "core/io/resource_loader.h"
 #include "core/io/resource_saver.h"
 #include "core/io/stream_peer_ssl.h"
-#include "core/message_queue.h"
+#include "core/object/class_db.h"
+#include "core/object/message_queue.h"
 #include "core/os/file_access.h"
 #include "core/os/keyboard.h"
 #include "core/os/os.h"
-#include "core/print_string.h"
-#include "core/project_settings.h"
-#include "core/translation.h"
+#include "core/string/print_string.h"
+#include "core/string/translation.h"
 #include "core/version.h"
 #include "main/main.h"
 #include "scene/gui/center_container.h"
@@ -2705,10 +2705,14 @@ void EditorNode::_screenshot(bool p_use_utc) {
 }
 
 void EditorNode::_save_screenshot(NodePath p_path) {
-	SubViewport *viewport = Object::cast_to<SubViewport>(EditorInterface::get_singleton()->get_editor_viewport()->get_viewport());
-	viewport->set_clear_mode(SubViewport::CLEAR_MODE_ONLY_NEXT_FRAME);
-	Ref<Image> img = viewport->get_texture()->get_data();
-	viewport->set_clear_mode(SubViewport::CLEAR_MODE_ALWAYS);
+	Control *editor_viewport = EditorInterface::get_singleton()->get_editor_viewport();
+	ERR_FAIL_COND_MSG(!editor_viewport, "Cannot get editor viewport.");
+	Viewport *viewport = editor_viewport->get_viewport();
+	ERR_FAIL_COND_MSG(!viewport, "Cannot get editor viewport.");
+	Ref<ViewportTexture> texture = viewport->get_texture();
+	ERR_FAIL_COND_MSG(texture.is_null(), "Cannot get editor viewport texture.");
+	Ref<Image> img = texture->get_data();
+	ERR_FAIL_COND_MSG(img.is_null(), "Cannot get editor viewport texture image.");
 	Error error = img->save_png(p_path);
 	ERR_FAIL_COND_MSG(error != OK, "Cannot save screenshot to file '" + p_path + "'.");
 }