Fix crash w/ scripts run from the command line that load files

Also refactor context backup/restore into functions.

1 files changed, 60 insertions, 23 deletions

diff --git a/source/creator/creator.c b/source/creator/creator.c
index ca8a9f8760d..797db2dfda1 100644
--- a/source/creator/creator.c
+++ b/source/creator/creator.c
@@ -1335,33 +1335,70 @@ static int set_skip_frame(int argc, const char **argv, void *data)
 	}
 }
 
-/* macro for ugly context setup/reset */
+
 #ifdef WITH_PYTHON
-#define BPY_CTX_SETUP(_cmd)                                                   \
-	{                                                                         \
-		wmWindowManager *wm = CTX_wm_manager(C);                              \
-		Scene *scene_prev = CTX_data_scene(C);                                \
-		wmWindow *win_prev;                                                   \
-		const bool has_win = !BLI_listbase_is_empty(&wm->windows);            \
-		if (has_win) {                                                        \
-			win_prev = CTX_wm_window(C);                                      \
-			CTX_wm_window_set(C, wm->windows.first);                          \
-		}                                                                     \
-		else {                                                                \
-			fprintf(stderr, "Python script \"%s\" "                           \
-			        "running with missing context data.\n", argv[1]);         \
-		}                                                                     \
-		{                                                                     \
-			_cmd;                                                             \
-		}                                                                     \
-		if (has_win) {                                                        \
-			CTX_wm_window_set(C, win_prev);                                   \
-		}                                                                     \
-		CTX_data_scene_set(C, scene_prev);                                    \
-	} (void)0                                                                 \
+
+/** \name Utilities Python Context Macro (#BPY_CTX_SETUP)
+ * \{ */
+struct BlendePyContextStore {
+	wmWindowManager *wm;
+	Scene *scene;
+	wmWindow *win;
+	bool has_win;
+};
+
+static void blender_py_context_backup(
+        bContext *C, struct BlendePyContextStore *c_py,
+        const char *script_id)
+{
+	c_py->wm = CTX_wm_manager(C);
+	c_py->scene = CTX_data_scene(C);
+	c_py->has_win = !BLI_listbase_is_empty(&c_py->wm->windows);
+	if (c_py->has_win) {
+		c_py->win = CTX_wm_window(C);
+		CTX_wm_window_set(C, c_py->wm->windows.first);
+	}
+	else {
+		c_py->win = NULL;
+		fprintf(stderr, "Python script \"%s\" "
+		        "running with missing context data.\n", script_id);
+	}
+}
+
+static void blender_py_context_restore(
+        bContext *C, struct BlendePyContextStore *c_py)
+{
+	/* script may load a file, check old data is valid before using */
+	if (c_py->has_win) {
+		if ((c_py->win == NULL) ||
+		    ((BLI_findindex(&G.main->wm, c_py->wm) != -1) &&
+		     (BLI_findindex(&c_py->wm->windows, c_py->win) != -1)))
+		{
+			CTX_wm_window_set(C, c_py->win);
+		}
+	}
+
+	if ((c_py->scene == NULL) ||
+	    BLI_findindex(&G.main->scene, c_py->scene) != -1)
+	{
+		CTX_data_scene_set(C, c_py->scene);
+	}
+}
+
+/* macro for context setup/reset */
+#define BPY_CTX_SETUP(_cmd) \
+	{ \
+		struct BlendePyContextStore py_c; \
+		blender_py_context_backup(C, &py_c, argv[1]); \
+		{ _cmd; } \
+		blender_py_context_restore(C, &py_c); \
+	} ((void)0)
 
 #endif /* WITH_PYTHON */
 
+/** \} */
+
+
 static int run_python_file(int argc, const char **argv, void *data)
 {
 #ifdef WITH_PYTHON