Enable adding scripts with redeemers for withdrawals

pycardano/txbuilder.py

@@ -138,6 +138,10 @@ class TransactionBuilder:
         init=False, default_factory=lambda: []
     )
 
+    _withdrawal_script_to_redeemers: List[
+        Tuple[ScriptType, Optional[Redeemer]]
+    ] = field(init=False, default_factory=lambda: [])
+
     _inputs_to_scripts: Dict[UTxO, ScriptType] = field(
         init=False, default_factory=lambda: {}
     )
@@ -291,6 +295,40 @@ def add_minting_script(
             self._minting_script_to_redeemers.append((script, redeemer))
         return self
 
+    def add_withdrawal_script(
+        self,
+        script: Union[UTxO, NativeScript, PlutusV1Script, PlutusV2Script],
+        redeemer: Optional[Redeemer] = None,
+    ) -> TransactionBuilder:
+        """Add a withdrawal script along with its redeemer to this transaction.
+
+        Args:
+            script (Union[UTxO, PlutusV1Script, PlutusV2Script]): A plutus script.
+            redeemer (Optional[Redeemer]): A plutus redeemer to unlock the UTxO.
+
+        Returns:
+            TransactionBuilder: Current transaction builder.
+        """
+        if redeemer:
+            if redeemer.tag is not None and redeemer.tag != RedeemerTag.REWARD:
+                raise InvalidArgumentException(
+                    f"Expect the redeemer tag's type to be {RedeemerTag.REWARD}, "
+                    f"but got {redeemer.tag} instead."
+                )
+            redeemer.tag = RedeemerTag.REWARD
+            self._consolidate_redeemer(redeemer)
+
+        if isinstance(script, UTxO):
+            assert script.output.script is not None
+            self._withdrawal_script_to_redeemers.append(
+                (script.output.script, redeemer)
+            )
+            self.reference_inputs.add(script)
+            self._reference_scripts.append(script.output.script)
+        else:
+            self._withdrawal_script_to_redeemers.append((script, redeemer))
+        return self
+
     def add_input_address(self, address: Union[Address, str]) -> TransactionBuilder:
         """Add an address to transaction's input address.
         Unlike :meth:`add_input`, which deterministically adds a UTxO to the transaction's inputs, `add_input_address`
@@ -376,6 +414,9 @@ def all_scripts(self) -> List[ScriptType]:
         for s, _ in self._minting_script_to_redeemers:
             scripts[script_hash(s)] = s
 
+        for s, _ in self._withdrawal_script_to_redeemers:
+            scripts[script_hash(s)] = s
+
         return list(scripts.values())
 
     @property
@@ -397,9 +438,11 @@ def datums(self) -> Dict[DatumHash, Datum]:
 
     @property
     def redeemers(self) -> List[Redeemer]:
-        return [r for r in self._inputs_to_redeemers.values() if r is not None] + [
-            r for _, r in self._minting_script_to_redeemers if r is not None
-        ]
+        return (
+            [r for r in self._inputs_to_redeemers.values() if r is not None]
+            + [r for _, r in self._minting_script_to_redeemers if r is not None]
+            + [r for _, r in self._withdrawal_script_to_redeemers if r is not None]
+        )
 
     @property
     def script_data_hash(self) -> Optional[ScriptDataHash]:
@@ -735,26 +778,31 @@ def _set_redeemer_index(self):
             sorted_mint_policies = sorted(self.mint.keys(), key=lambda x: x.to_cbor())
         else:
             sorted_mint_policies = []
+        if self.withdrawals:
+            sorted_withdrawals = sorted(self.withdrawals.keys())
+        else:
+            sorted_withdrawals = []
 
         for i, utxo in enumerate(self.inputs):
             if (
                 utxo in self._inputs_to_redeemers
                 and self._inputs_to_redeemers[utxo].tag == RedeemerTag.SPEND
             ):
                 self._inputs_to_redeemers[utxo].index = i
-            elif (
-                utxo in self._inputs_to_redeemers
-                and self._inputs_to_redeemers[utxo].tag == RedeemerTag.MINT
-            ):
-                redeemer = self._inputs_to_redeemers[utxo]
-                redeemer.index = sorted_mint_policies.index(
-                    script_hash(self._inputs_to_scripts[utxo])
-                )
 
         for script, redeemer in self._minting_script_to_redeemers:
             if redeemer is not None:
                 redeemer.index = sorted_mint_policies.index(script_hash(script))
 
+        for script, redeemer in self._withdrawal_script_to_redeemers:
+            if redeemer is not None:
+                script_staking_credential = Address(
+                    staking_part=script_hash(script), network=self.context.network
+                )
+                redeemer.index = sorted_withdrawals.index(
+                    script_staking_credential.to_primitive()
+                )
+
         self.redeemers.sort(key=lambda r: r.index)
 
     def _build_tx_body(self) -> TransactionBody:
@@ -1203,7 +1251,10 @@ def _update_execution_units(
                 assert (
                     r.tag is not None
                 ), "Expected tag of redeemer to be set, but found None"
-                key = f"{r.tag.name.lower()}:{r.index}"
+                tagname = (
+                    r.tag.name.lower() if r.tag != RedeemerTag.REWARD else "withdrawal"
+                )
+                key = f"{tagname}:{r.index}"
                 if (
                     key not in estimated_execution_units
                     or estimated_execution_units[key] is None