diff --git a/pycardano/txbuilder.py b/pycardano/txbuilder.py index e3e68471..9a1c691b 100644 --- a/pycardano/txbuilder.py +++ b/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,6 +778,10 @@ 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 ( @@ -742,19 +789,20 @@ def _set_redeemer_index(self): 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