From 815337a7dafcfe1504c976fb68f6f988749389a2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Fri, 31 Jan 2025 15:16:10 +0100
Subject: [PATCH] Anim: add `ensure` parameter to `KeyframeStrip.channels(...)`
 RNA function

Add `ensure` boolean parameter to the `ActionKeyframeStrip.channels()` RNA
function. Passing `ensure=True` will ensure the channelbag for the given
action slot handle exists. This makes it more straight-forward to create
F-Curves for a slot regardless of whether there already was a channelbag
for those F-Cuves:

```python
strip = action.layers[0].strips[0]
slot = action.slots[0]

# Old:
channelbag = strip.channels(slot.handle)
if not channelbag:
    channelbag = strip.channelbags.new(action_slot)
ob_loc_x = channelbag.fcurves.new('location', index=0)

# New:
channelbag = strip.channels(slot.handle, ensure=True)
ob_loc_x = channelbag.fcurves.new('location', index=0)
```

Pull Request: https://projects.blender.org/blender/blender/pulls/133678
---
 source/blender/animrig/ANIM_action.hh        |  1 +
 source/blender/animrig/intern/action.cc      |  9 +++++++--
 source/blender/makesrna/intern/rna_action.cc | 12 +++++++++++-
 tests/python/bl_animation_action.py          |  4 ++++
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/source/blender/animrig/ANIM_action.hh b/source/blender/animrig/ANIM_action.hh
index 8fad04acc9a..7adf612c6c1 100644
--- a/source/blender/animrig/ANIM_action.hh
+++ b/source/blender/animrig/ANIM_action.hh
@@ -971,6 +971,7 @@ class StripKeyframeData : public ::ActionStripKeyframeData {
    * Find the channelbag for the given slot, or if none exists, create it.
    */
   Channelbag &channelbag_for_slot_ensure(const Slot &slot);
+  Channelbag &channelbag_for_slot_ensure(slot_handle_t slot_handle);
 
   /**
    * Remove the given channelbag from this strip data.
diff --git a/source/blender/animrig/intern/action.cc b/source/blender/animrig/intern/action.cc
index 61640414da1..d28c2e5eb54 100644
--- a/source/blender/animrig/intern/action.cc
+++ b/source/blender/animrig/intern/action.cc
@@ -1774,11 +1774,16 @@ Channelbag &StripKeyframeData::channelbag_for_slot_add(const slot_handle_t slot_
 
 Channelbag &StripKeyframeData::channelbag_for_slot_ensure(const Slot &slot)
 {
-  Channelbag *channelbag = this->channelbag_for_slot(slot);
+  return this->channelbag_for_slot_ensure(slot.handle);
+}
+
+Channelbag &StripKeyframeData::channelbag_for_slot_ensure(const slot_handle_t slot_handle)
+{
+  Channelbag *channelbag = this->channelbag_for_slot(slot_handle);
   if (channelbag != nullptr) {
     return *channelbag;
   }
-  return this->channelbag_for_slot_add(slot);
+  return this->channelbag_for_slot_add(slot_handle);
 }
 
 static void channelbag_ptr_destructor(ActionChannelbag **dna_channelbag_ptr)
diff --git a/source/blender/makesrna/intern/rna_action.cc b/source/blender/makesrna/intern/rna_action.cc
index 21a8f90bda4..2137733aa57 100644
--- a/source/blender/makesrna/intern/rna_action.cc
+++ b/source/blender/makesrna/intern/rna_action.cc
@@ -721,10 +721,15 @@ static void rna_Channelbag_group_remove(ActionChannelbag *dna_channelbag,
 
 static ActionChannelbag *rna_ActionStrip_channels(ID *dna_action_id,
                                                   ActionStrip *self,
-                                                  const animrig::slot_handle_t slot_handle)
+                                                  const animrig::slot_handle_t slot_handle,
+                                                  const bool ensure)
 {
   animrig::Action &action = reinterpret_cast<bAction *>(dna_action_id)->wrap();
   animrig::StripKeyframeData &strip_data = self->wrap().data<animrig::StripKeyframeData>(action);
+
+  if (ensure) {
+    return &strip_data.channelbag_for_slot_ensure(slot_handle);
+  }
   return strip_data.channelbag_for_slot(slot_handle);
 }
 
@@ -2233,6 +2238,11 @@ static void rna_def_action_keyframe_strip(BlenderRNA *brna)
                        0,
                        INT_MAX);
     RNA_def_parameter_flags(parm, PropertyFlag(0), PARM_REQUIRED);
+    RNA_def_boolean(func,
+                    "ensure",
+                    false,
+                    "Create if necessary",
+                    "Ensure the channelbag exists for this slot handle, creating it if necessary");
     parm = RNA_def_pointer(func, "channels", "ActionChannelbag", "Channels", "");
     RNA_def_function_return(func, parm);
 
diff --git a/tests/python/bl_animation_action.py b/tests/python/bl_animation_action.py
index 3ba20a59397..cf2d698bf91 100644
--- a/tests/python/bl_animation_action.py
+++ b/tests/python/bl_animation_action.py
@@ -395,6 +395,10 @@ class ChannelbagsTest(unittest.TestCase):
         self.strip.channelbags.remove(channelbag)
         self.assertEqual([], list(self.strip.channelbags))
 
+    def test_ensure_channelbag(self):
+        channelbag = self.strip.channels(self.slot.handle, ensure=True)
+        self.assertEqual([channelbag], list(self.strip.channelbags))
+
     def test_create_remove_fcurves(self):
         channelbag = self.strip.channelbags.new(self.slot)