PyAPI: temporary context override support

Support a way to temporarily override the context from Python.

- Added method `Context.temp_override` context manager.
- Special support for windowing variables "window", "area" and "region",
  other context members such as "active_object".
- Nesting context overrides is supported.
- Previous windowing members are restored when the context exists unless
  they have been removed.
- Overriding context members by passing a dictionary into operators in
  `bpy.ops` has been deprecated and warns when used.

This allows the window in a newly loaded file to be used, see: T92464

Reviewed by: mont29

Ref D13126

8 files changed, 404 insertions, 2 deletions

diff --git a/doc/python_api/examples/bpy.types.Context.temp_override.1.py b/doc/python_api/examples/bpy.types.Context.temp_override.1.py
new file mode 100644
index 00000000000..68f0eef93c3
--- /dev/null
+++ b/doc/python_api/examples/bpy.types.Context.temp_override.1.py
@@ -0,0 +1,19 @@
+"""
+Overriding the context can be used to temporarily activate another ``window`` / ``area`` & ``region``,
+as well as other members such as the ``active_object`` or ``bone``.
+
+Notes:
+
+- When overriding window, area and regions: the arguments must be consistent,
+  so any region argument that's passed in must be contained by the current area or the area passed in.
+  The same goes for the area needing to be contained in the current window.
+
+- Temporary context overrides may be nested, when this is done, members will be added to the existing overrides.
+
+- Context members are restored outside the scope of the context.
+  The only exception to this is when the data is no longer available.
+
+  In the event windowing data was removed (for example), the state of the context is left as-is.
+  While this isn't likely to happen, explicit window operation such as closing windows or loading a new file
+  remove the windowing data that was set before the temporary context was created.
+"""
diff --git a/doc/python_api/examples/bpy.types.Context.temp_override.2.py b/doc/python_api/examples/bpy.types.Context.temp_override.2.py
new file mode 100644
index 00000000000..ce3e1594baa
--- /dev/null
+++ b/doc/python_api/examples/bpy.types.Context.temp_override.2.py
@@ -0,0 +1,15 @@
+"""
+Overriding the context can be useful to set the context after loading files
+(which would otherwise by None). For example:
+"""
+
+import bpy
+from bpy import context
+
+# Reload the current file and select all.
+bpy.ops.wm.open_mainfile(filepath=bpy.data.filepath)
+window = context.window_manager.windows[0]
+with context.temp_override(window=window):
+    bpy.ops.mesh.primitive_uv_sphere_add()
+    # The context override is needed so it's possible to set edit-mode.
+    bpy.ops.object.mode_set(mode='EDIT')
diff --git a/doc/python_api/examples/bpy.types.Context.temp_override.3.py b/doc/python_api/examples/bpy.types.Context.temp_override.3.py
new file mode 100644
index 00000000000..e670bb7bafa
--- /dev/null
+++ b/doc/python_api/examples/bpy.types.Context.temp_override.3.py
@@ -0,0 +1,16 @@
+"""
+This example shows how it's possible to add an object to the scene in another window.
+"""
+import bpy
+from bpy import context
+
+win_active = context.window
+win_other = None
+for win_iter in context.window_manager.windows:
+    if win_iter != win_active:
+        win_other = win_iter
+        break
+
+# Add cube in the other window.
+with context.temp_override(window=win_other):
+    bpy.ops.mesh.primitive_cube_add()
diff --git a/source/blender/python/intern/CMakeLists.txt b/source/blender/python/intern/CMakeLists.txt
index 86dc5800b67..e4e198ab812 100644
--- a/source/blender/python/intern/CMakeLists.txt
+++ b/source/blender/python/intern/CMakeLists.txt
@@ -60,6 +60,7 @@ set(SRC
   bpy_rna_anim.c
   bpy_rna_array.c
   bpy_rna_callback.c
+  bpy_rna_context.c
   bpy_rna_data.c
   bpy_rna_driver.c
   bpy_rna_gizmo.c
@@ -101,6 +102,7 @@ set(SRC
   bpy_rna.h
   bpy_rna_anim.h
   bpy_rna_callback.h
+  bpy_rna_context.h
   bpy_rna_data.h
   bpy_rna_driver.h
   bpy_rna_gizmo.h
diff --git a/source/blender/python/intern/bpy_operator.c b/source/blender/python/intern/bpy_operator.c
index db0067fc18e..0cfe6dab2f5 100644
--- a/source/blender/python/intern/bpy_operator.c
+++ b/source/blender/python/intern/bpy_operator.c
@@ -60,6 +60,18 @@ static wmOperatorType *ot_lookup_from_py_string(PyObject *value, const char *py_
   return ot;
 }
 
+static void op_context_override_deprecated_warning(void)
+{
+  if (PyErr_WarnEx(PyExc_DeprecationWarning,
+                   "Passing in context overrides is deprecated in favor of "
+                   "Context.temp_override(..)",
+                   1) < 0) {
+    /* The function has no return value, the exception cannot
+     * be reported to the caller, so just log it. */
+    PyErr_WriteUnraisable(NULL);
+  }
+}
+
 static PyObject *pyop_poll(PyObject *UNUSED(self), PyObject *args)
 {
   wmOperatorType *ot;
@@ -113,7 +125,10 @@ static PyObject *pyop_poll(PyObject *UNUSED(self), PyObject *args)
   if (ELEM(context_dict, NULL, Py_None)) {
     context_dict = NULL;
   }
-  else if (!PyDict_Check(context_dict)) {
+  else if (PyDict_Check(context_dict)) {
+    op_context_override_deprecated_warning();
+  }
+  else {
     PyErr_Format(PyExc_TypeError,
                  "Calling operator \"bpy.ops.%s.poll\" error, "
                  "custom context expected a dict or None, got a %.200s",
@@ -218,7 +233,10 @@ static PyObject *pyop_call(PyObject *UNUSED(self), PyObject *args)
   if (ELEM(context_dict, NULL, Py_None)) {
     context_dict = NULL;
   }
-  else if (!PyDict_Check(context_dict)) {
+  else if (PyDict_Check(context_dict)) {
+    op_context_override_deprecated_warning();
+  }
+  else {
     PyErr_Format(PyExc_TypeError,
                  "Calling operator \"bpy.ops.%s\" error, "
                  "custom context expected a dict or None, got a %.200s",
diff --git a/source/blender/python/intern/bpy_rna_context.c b/source/blender/python/intern/bpy_rna_context.c
new file mode 100644
index 00000000000..085a8323cc1
--- /dev/null
+++ b/source/blender/python/intern/bpy_rna_context.c
@@ -0,0 +1,299 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+/** \file
+ * \ingroup pythonintern
+ *
+ * This file adds some helper methods to the context, that cannot fit well in RNA itself.
+ */
+
+#include <Python.h>
+
+#include "BLI_listbase.h"
+#include "BLI_utildefines.h"
+
+#include "BKE_context.h"
+
+#include "WM_api.h"
+#include "WM_types.h"
+
+#include "bpy_rna_context.h"
+
+#include "RNA_access.h"
+#include "RNA_prototypes.h"
+
+#include "bpy_rna.h"
+
+/* -------------------------------------------------------------------- */
+/** \name Temporary Context Override (Python Context Manager)
+ * \{ */
+
+typedef struct ContextStore {
+  wmWindow *win;
+  bool win_is_set;
+  ScrArea *area;
+  bool area_is_set;
+  ARegion *region;
+  bool region_is_set;
+} ContextStore;
+
+typedef struct BPyContextTempOverride {
+  PyObject_HEAD /* Required Python macro. */
+  bContext *context;
+
+  ContextStore ctx_init;
+  ContextStore ctx_temp;
+  /** Bypass Python overrides set when calling an operator from Python. */
+  struct bContext_PyState py_state;
+  /**
+   * This dictionary is used to store members that don't have special handling,
+   * see: #bpy_context_temp_override_extract_known_args,
+   * these will then be accessed via #BPY_context_member_get.
+   *
+   * This also supports nested *stacking*, so a nested temp-context-overrides
+   * will overlay the new members on the old members (instead of ignoring them).
+   */
+  PyObject *py_state_context_dict;
+} BPyContextTempOverride;
+
+static void bpy_rna_context_temp_override__tp_dealloc(BPyContextTempOverride *self)
+{
+  PyObject_DEL(self);
+}
+
+static PyObject *bpy_rna_context_temp_override_enter(BPyContextTempOverride *self)
+{
+  bContext *C = self->context;
+
+  CTX_py_state_push(C, &self->py_state, self->py_state_context_dict);
+
+  self->ctx_init.win = CTX_wm_window(C);
+  self->ctx_init.win_is_set = (self->ctx_init.win != self->ctx_temp.win);
+  self->ctx_init.area = CTX_wm_area(C);
+  self->ctx_init.area_is_set = (self->ctx_init.area != self->ctx_temp.area);
+  self->ctx_init.region = CTX_wm_region(C);
+  self->ctx_init.region_is_set = (self->ctx_init.region != self->ctx_temp.region);
+
+  wmWindow *win = self->ctx_temp.win_is_set ? self->ctx_temp.win : self->ctx_init.win;
+  bScreen *screen = win ? WM_window_get_active_screen(win) : NULL;
+  ScrArea *area = self->ctx_temp.area_is_set ? self->ctx_temp.area : self->ctx_init.area;
+  ARegion *region = self->ctx_temp.region_is_set ? self->ctx_temp.region : self->ctx_init.region;
+
+  /* Sanity check, the region is in the screen/area. */
+  if (self->ctx_temp.region_is_set && (region != NULL)) {
+    if (area == NULL) {
+      PyErr_SetString(PyExc_TypeError, "Region set with NULL area");
+      return NULL;
+    }
+    if ((screen && BLI_findindex(&screen->regionbase, region) == -1) &&
+        (BLI_findindex(&area->regionbase, region) == -1)) {
+      PyErr_SetString(PyExc_TypeError, "Region not found in area");
+      return NULL;
+    }
+  }
+
+  if (self->ctx_temp.area_is_set && (area != NULL)) {
+    if (screen == NULL) {
+      PyErr_SetString(PyExc_TypeError, "Area set with NULL screen");
+      return NULL;
+    }
+    if (BLI_findindex(&screen->areabase, area) == -1) {
+      PyErr_SetString(PyExc_TypeError, "Area not found in screen");
+      return NULL;
+    }
+  }
+
+  if (self->ctx_temp.win_is_set) {
+    CTX_wm_window_set(C, self->ctx_temp.win);
+  }
+  if (self->ctx_temp.area_is_set) {
+    CTX_wm_area_set(C, self->ctx_temp.area);
+  }
+  if (self->ctx_temp.region_is_set) {
+    CTX_wm_region_set(C, self->ctx_temp.region);
+  }
+
+  Py_RETURN_NONE;
+}
+
+static PyObject *bpy_rna_context_temp_override_exit(BPyContextTempOverride *self,
+                                                    PyObject *UNUSED(args))
+{
+  bContext *C = self->context;
+
+  /* Special case where the window is expected to be freed on file-read,
+   * in this case the window should not be restored, see: T92818. */
+  bool do_restore = true;
+  if (self->ctx_init.win) {
+    wmWindowManager *wm = CTX_wm_manager(C);
+    if (BLI_findindex(&wm->windows, self->ctx_init.win) == -1) {
+      CTX_wm_window_set(C, NULL);
+      do_restore = false;
+    }
+  }
+
+  if (do_restore) {
+    if (self->ctx_init.win_is_set) {
+      CTX_wm_window_set(C, self->ctx_init.win);
+    }
+    if (self->ctx_init.area_is_set) {
+      CTX_wm_area_set(C, self->ctx_init.area);
+    }
+    if (self->ctx_init.region_is_set) {
+      CTX_wm_region_set(C, self->ctx_init.region);
+    }
+  }
+
+  CTX_py_state_pop(C, &self->py_state);
+  Py_CLEAR(self->py_state_context_dict);
+
+  Py_RETURN_NONE;
+}
+
+static PyMethodDef bpy_rna_context_temp_override__tp_methods[] = {
+    {"__enter__", (PyCFunction)bpy_rna_context_temp_override_enter, METH_NOARGS},
+    {"__exit__", (PyCFunction)bpy_rna_context_temp_override_exit, METH_VARARGS},
+    {NULL},
+};
+
+static PyTypeObject BPyContextTempOverride_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0).tp_name = "ContextTempOverride",
+    .tp_basicsize = sizeof(BPyContextTempOverride),
+    .tp_dealloc = (destructor)bpy_rna_context_temp_override__tp_dealloc,
+    .tp_flags = Py_TPFLAGS_DEFAULT,
+    .tp_methods = bpy_rna_context_temp_override__tp_methods,
+};
+
+/** \} */
+
+/* -------------------------------------------------------------------- */
+/** \name Context Temporary Override Method
+ * \{ */
+
+static PyObject *bpy_context_temp_override_extract_known_args(const char *const *kwds_static,
+                                                              PyObject *kwds)
+{
+  PyObject *sentinel = Py_Ellipsis;
+  PyObject *kwds_parse = PyDict_New();
+  for (int i = 0; kwds_static[i]; i++) {
+    PyObject *key = PyUnicode_FromString(kwds_static[i]);
+    PyObject *val = _PyDict_Pop(kwds, key, sentinel);
+    if (val != sentinel) {
+      if (PyDict_SetItem(kwds_parse, key, val) == -1) {
+        BLI_assert_unreachable();
+      }
+    }
+    Py_DECREF(key);
+    Py_DECREF(val);
+  }
+  return kwds_parse;
+}
+
+PyDoc_STRVAR(bpy_context_temp_override_doc,
+             ".. method:: temp_override(window, area, region, **keywords)\n"
+             "\n"
+             "   Context manager to temporarily override members in the context.\n"
+             "\n"
+             "   :arg window: Window override or None.\n"
+             "   :type window: :class:`bpy.types.Window`\n"
+             "   :arg area: Area override or None.\n"
+             "   :type area: :class:`bpy.types.Area`\n"
+             "   :arg region: Region override or None.\n"
+             "   :type region: :class:`bpy.types.Region`\n"
+             "   :arg keywords: Additional keywords override context members.\n"
+             "   :return: The context manager .\n"
+             "   :rtype: context manager\n");
+static PyObject *bpy_context_temp_override(PyObject *self, PyObject *args, PyObject *kwds)
+{
+  const PointerRNA *context_ptr = pyrna_struct_as_ptr(self, &RNA_Context);
+  if (context_ptr == NULL) {
+    return NULL;
+  }
+  /* Needed because the keywords copied into `kwds_parse` could contain anything.
+   * As the types of keys aren't checked. */
+  if (!PyArg_ValidateKeywordArguments(kwds)) {
+    return NULL;
+  }
+
+  struct {
+    struct BPy_StructRNA_Parse window;
+    struct BPy_StructRNA_Parse area;
+    struct BPy_StructRNA_Parse region;
+  } params = {
+      .window = {.type = &RNA_Window},
+      .area = {.type = &RNA_Area},
+      .region = {.type = &RNA_Region},
+  };
+
+  static const char *const _keywords[] = {"window", "area", "region", NULL};
+  static _PyArg_Parser _parser = {
+      "|$" /* Optional, keyword only arguments. */
+      "O&" /* `window` */
+      "O&" /* `area` */
+      "O&" /* `region` */
+      ":temp_override",
+      _keywords,
+      0,
+  };
+  /* Parse known keywords, the remaining keywords are set using #CTX_py_state_push. */
+  kwds = PyDict_Copy(kwds);
+  {
+    PyObject *kwds_parse = bpy_context_temp_override_extract_known_args(_keywords, kwds);
+    const int parse_result = _PyArg_ParseTupleAndKeywordsFast(args,
+                                                              kwds_parse,
+                                                              &_parser,
+                                                              pyrna_struct_as_ptr_or_null_parse,
+                                                              &params.window,
+                                                              pyrna_struct_as_ptr_or_null_parse,
+                                                              &params.area,
+                                                              pyrna_struct_as_ptr_or_null_parse,
+                                                              &params.region);
+    Py_DECREF(kwds_parse);
+    if (parse_result == -1) {
+      Py_DECREF(kwds);
+      return NULL;
+    }
+  }
+
+  bContext *C = context_ptr->data;
+  {
+    /* Merge existing keys that don't exist in the keywords passed in.
+     * This makes it possible to nest context overrides. */
+    PyObject *context_dict_current = CTX_py_dict_get(C);
+    if (context_dict_current != NULL) {
+      PyDict_Merge(kwds, context_dict_current, 0);
+    }
+  }
+
+  ContextStore ctx_temp = {NULL};
+  if (params.window.ptr != NULL) {
+    ctx_temp.win = params.window.ptr->data;
+    ctx_temp.win_is_set = true;
+  }
+  if (params.area.ptr != NULL) {
+    ctx_temp.area = params.area.ptr->data;
+    ctx_temp.area_is_set = true;
+  }
+
+  if (params.region.ptr != NULL) {
+    ctx_temp.region = params.region.ptr->data;
+    ctx_temp.region_is_set = true;
+  }
+
+  BPyContextTempOverride *ret = PyObject_New(BPyContextTempOverride, &BPyContextTempOverride_Type);
+  ret->context = C;
+  ret->ctx_temp = ctx_temp;
+  memset(&ret->ctx_init, 0, sizeof(ret->ctx_init));
+
+  ret->py_state_context_dict = kwds;
+
+  return (PyObject *)ret;
+}
+
+/** \} */
+
+PyMethodDef BPY_rna_context_temp_override_method_def = {
+    "temp_override",
+    (PyCFunction)bpy_context_temp_override,
+    METH_VARARGS | METH_KEYWORDS,
+    bpy_context_temp_override_doc,
+};
diff --git a/source/blender/python/intern/bpy_rna_context.h b/source/blender/python/intern/bpy_rna_context.h
new file mode 100644
index 00000000000..ddd328131e6
--- /dev/null
+++ b/source/blender/python/intern/bpy_rna_context.h
@@ -0,0 +1,17 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+/** \file
+ * \ingroup pythonintern
+ */
+
+#pragma once
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+extern PyMethodDef BPY_rna_context_temp_override_method_def;
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/source/blender/python/intern/bpy_rna_types_capi.c b/source/blender/python/intern/bpy_rna_types_capi.c
index a5299bc1616..376195ab845 100644
--- a/source/blender/python/intern/bpy_rna_types_capi.c
+++ b/source/blender/python/intern/bpy_rna_types_capi.c
@@ -22,6 +22,7 @@
 #include "bpy_library.h"
 #include "bpy_rna.h"
 #include "bpy_rna_callback.h"
+#include "bpy_rna_context.h"
 #include "bpy_rna_data.h"
 #include "bpy_rna_id_collection.h"
 #include "bpy_rna_text.h"
@@ -159,6 +160,17 @@ static struct PyGetSetDef pyrna_windowmanager_getset[] = {
 /** \} */
 
 /* -------------------------------------------------------------------- */
+/** \name Context Type
+ * \{ */
+
+static struct PyMethodDef pyrna_context_methods[] = {
+    {NULL, NULL, 0, NULL}, /* #BPY_rna_context_temp_override_method_def */
+    {NULL, NULL, 0, NULL},
+};
+
+/** \} */
+
+/* -------------------------------------------------------------------- */
 /** \name Space Type
  * \{ */
 
@@ -254,6 +266,10 @@ void BPY_rna_types_extend_capi(void)
   /* WindowManager */
   pyrna_struct_type_extend_capi(
       &RNA_WindowManager, pyrna_windowmanager_methods, pyrna_windowmanager_getset);
+
+  /* Context */
+  ARRAY_SET_ITEMS(pyrna_context_methods, BPY_rna_context_temp_override_method_def);
+  pyrna_struct_type_extend_capi(&RNA_Context, pyrna_context_methods, NULL);
 }
 
 /** \} */