4 files changed, 42 insertions, 4 deletions

diff --git a/source/blender/blenkernel/BKE_context.h b/source/blender/blenkernel/BKE_context.h
index 584568b4a5e..f441203d9d6 100644
--- a/source/blender/blenkernel/BKE_context.h
+++ b/source/blender/blenkernel/BKE_context.h
@@ -10,6 +10,8 @@
 /* XXX temporary, until AssetHandle is designed properly and queries can return a pointer to it. */
 #include "DNA_asset_types.h"
 
+#include "BLI_utildefines.h"
+
 #include "DNA_listBase.h"
 #include "DNA_object_enums.h"
 #include "RNA_types.h"
@@ -136,6 +138,9 @@ bContextStore *CTX_store_add(ListBase *contexts, const char *name, const Pointer
 bContextStore *CTX_store_add_all(ListBase *contexts, bContextStore *context);
 bContextStore *CTX_store_get(bContext *C);
 void CTX_store_set(bContext *C, bContextStore *store);
+const PointerRNA *CTX_store_ptr_lookup(const bContextStore *store,
+                                       const char *name,
+                                       const StructRNA *type CPP_ARG_DEFAULT(nullptr));
 bContextStore *CTX_store_copy(bContextStore *store);
 void CTX_store_free(bContextStore *store);
 void CTX_store_free_list(ListBase *contexts);
diff --git a/source/blender/blenkernel/intern/context.c b/source/blender/blenkernel/intern/context.c
index 63efbb99368..1f1a49ca030 100644
--- a/source/blender/blenkernel/intern/context.c
+++ b/source/blender/blenkernel/intern/context.c
@@ -190,6 +190,22 @@ void CTX_store_set(bContext *C, bContextStore *store)
   C->wm.store = store;
 }
 
+const PointerRNA *CTX_store_ptr_lookup(const bContextStore *store,
+                                       const char *name,
+                                       const StructRNA *type)
+{
+  bContextStoreEntry *entry = BLI_rfindstring(
+      &store->entries, name, offsetof(bContextStoreEntry, name));
+  if (!entry) {
+    return NULL;
+  }
+
+  if (type && !RNA_struct_is_a(entry->ptr.type, type)) {
+    return NULL;
+  }
+  return &entry->ptr;
+}
+
 bContextStore *CTX_store_copy(bContextStore *store)
 {
   bContextStore *ctx = MEM_dupallocN(store);
@@ -324,11 +340,10 @@ static eContextResult ctx_data_get(bContext *C, const char *member, bContextData
   if (done != 1 && recursion < 1 && C->wm.store) {
     C->data.recursion = 1;
 
-    bContextStoreEntry *entry = BLI_rfindstring(
-        &C->wm.store->entries, member, offsetof(bContextStoreEntry, name));
+    const PointerRNA *ptr = CTX_store_ptr_lookup(C->wm.store, member, NULL);
 
-    if (entry) {
-      result->ptr = entry->ptr;
+    if (ptr) {
+      result->ptr = *ptr;
       done = 1;
     }
   }
diff --git a/source/blender/editors/include/UI_interface.h b/source/blender/editors/include/UI_interface.h
index 3f58aab3e53..94453ef5fac 100644
--- a/source/blender/editors/include/UI_interface.h
+++ b/source/blender/editors/include/UI_interface.h
@@ -1360,6 +1360,13 @@ uiBut *uiDefIconTextButO_ptr(uiBlock *block,
 /* for passing inputs to ButO buttons */
 struct PointerRNA *UI_but_operator_ptr_get(uiBut *but);
 
+void UI_but_context_ptr_set(uiBlock *block,
+                            uiBut *but,
+                            const char *name,
+                            const struct PointerRNA *ptr);
+const struct PointerRNA *UI_but_context_ptr_get(const uiBut *but,
+                                                const char *name,
+                                                const StructRNA *type CPP_ARG_DEFAULT(nullptr));
 struct bContextStore *UI_but_context_get(const uiBut *but);
 
 void UI_but_unit_type_set(uiBut *but, int unit_type);
diff --git a/source/blender/editors/interface/interface.cc b/source/blender/editors/interface/interface.cc
index 0ecaead4194..9310ede31fa 100644
--- a/source/blender/editors/interface/interface.cc
+++ b/source/blender/editors/interface/interface.cc
@@ -5951,6 +5951,17 @@ PointerRNA *UI_but_operator_ptr_get(uiBut *but)
   return but->opptr;
 }
 
+void UI_but_context_ptr_set(uiBlock *block, uiBut *but, const char *name, const PointerRNA *ptr)
+{
+  but->context = CTX_store_add(&block->contexts, name, ptr);
+  but->context->used = true;
+}
+
+const PointerRNA *UI_but_context_ptr_get(const uiBut *but, const char *name, const StructRNA *type)
+{
+  return CTX_store_ptr_lookup(but->context, name, type);
+}
+
 bContextStore *UI_but_context_get(const uiBut *but)
 {
   return but->context;