Merge change into already existing output when possible (#38)

* Add a copy method to TransactionOutput

* Add change to TxOutput with matching address when possible

* Update Tx Output addition method

* Update outputs and fees of relevant tests

* Lint the changes

* Update the copy method of TransactionOutput to __copy__

* Add merge change as an optional parameter

* Restore original tests before merge_change was default

* Do not merge change if there are multiple change UTxOs

* Lint new changes

* Remove unused imports

* Add unit tests and minor fixes for change merging

Co-authored-by: Jarred <[email protected]>, Jerry <[email protected]>

Author: astrojarred

Makefile

@@ -27,8 +27,7 @@ help:
 	@python -c "$$PRINT_HELP_PYSCRIPT" < $(MAKEFILE_LIST)
 
 cov: ## check code coverage
-	poetry run coverage run --source pycardano -m pytest -n 4
-	poetry run coverage report -m
+	poetry run pytest -n 4 --cov pycardano
 
 cov-html: cov ## check code coverage and generate an html report
 	poetry run coverage html -d cov_html

pycardano/txbuilder.py

@@ -312,7 +312,9 @@ def _calc_change(
         # when there is only ADA left, simply use remaining coin value as change
         if not change.multi_asset:
             if change.coin < min_lovelace(change, self.context):
-                raise InsufficientUTxOBalanceException("Not enough ADA left for change")
+                raise InsufficientUTxOBalanceException(
+                    f"Not enough ADA left for change: {change.coin} but needs {min_lovelace(change, self.context)}"
+                )
             lovelace_change = change.coin
             change_output_arr.append(TransactionOutput(address, lovelace_change))
 
@@ -347,17 +349,42 @@ def _calc_change(
         return change_output_arr
 
     def _add_change_and_fee(
-        self, change_address: Optional[Address]
+        self,
+        change_address: Optional[Address],
+        merge_change: Optional[bool] = False,
     ) -> TransactionBuilder:
-        original_outputs = self.outputs[:]
+        original_outputs = deepcopy(self.outputs)
+        change_output_index = None
+
+        def _merge_changes(changes):
+            if change_output_index is not None and len(changes) == 1:
+                # Add the leftover change to the TransactionOutput containing the change address
+                self._outputs[change_output_index].amount = (
+                    changes[0].amount + self._outputs[change_output_index].amount
+                )
+                # if we enforce that TransactionOutputs must use Values for `amount`, we can use += here
+
+            else:
+                self._outputs += changes
 
         if change_address:
+
+            if merge_change:
+
+                for idx, output in enumerate(original_outputs):
+
+                    # Find any transaction outputs which already contain the change address
+                    if change_address == output.address:
+                        if change_output_index is None or output.lovelace == 0:
+                            change_output_index = idx
+
             # Set fee to max
             self.fee = self._estimate_fee()
             changes = self._calc_change(
                 self.fee, self.inputs, self.outputs, change_address, precise_fee=True
             )
-            self._outputs += changes
+
+            _merge_changes(changes)
 
         # With changes included, we can estimate the fee more precisely
         self.fee = self._estimate_fee()
@@ -367,7 +394,8 @@ def _add_change_and_fee(
             changes = self._calc_change(
                 self.fee, self.inputs, self.outputs, change_address, precise_fee=True
             )
-            self._outputs += changes
+
+            _merge_changes(changes)
 
         return self
 
@@ -649,12 +677,18 @@ def _estimate_fee(self):
 
         return estimated_fee
 
-    def build(self, change_address: Optional[Address] = None) -> TransactionBody:
+    def build(
+        self,
+        change_address: Optional[Address] = None,
+        merge_change: Optional[bool] = False,
+    ) -> TransactionBody:
         """Build a transaction body from all constraints set through the builder.
 
         Args:
             change_address (Optional[Address]): Address to which changes will be returned. If not provided, the
                 transaction body will likely be unbalanced (sum of inputs is greater than the sum of outputs).
+            merge_change (Optional[bool]): If the change address match one of the transaction output, the change amount
+                will be directly added to that transaction output, instead of being added as a separate output.
 
         Returns:
             TransactionBody: A transaction body.
@@ -773,7 +807,7 @@ def build(self, change_address: Optional[Address] = None) -> TransactionBody:
 
         self._set_redeemer_index()
 
-        self._add_change_and_fee(change_address)
+        self._add_change_and_fee(change_address, merge_change=merge_change)
 
         tx_body = self._build_tx_body()
 
@@ -820,6 +854,7 @@ def build_and_sign(
         self,
         signing_keys: List[Union[SigningKey, ExtendedSigningKey]],
         change_address: Optional[Address] = None,
+        merge_change: Optional[bool] = False,
     ) -> Transaction:
         """Build a transaction body from all constraints set through the builder and sign the transaction with
         provided signing keys.
@@ -829,12 +864,14 @@ def build_and_sign(
                 sign the transaction.
             change_address (Optional[Address]): Address to which changes will be returned. If not provided, the
                 transaction body will likely be unbalanced (sum of inputs is greater than the sum of outputs).
+            merge_change (Optional[bool]): If the change address match one of the transaction output, the change amount
+                will be directly added to that transaction output, instead of being added as a separate output.
 
         Returns:
             Transaction: A signed transaction.
         """
 
-        tx_body = self.build(change_address=change_address)
+        tx_body = self.build(change_address=change_address, merge_change=merge_change)
         witness_set = self.build_witness_set()
         witness_set.vkey_witnesses = []
 

test/pycardano/test_txbuilder.py

@@ -703,3 +703,116 @@ def test_tx_builder_withdrawal(chain_context):
     }
 
     assert expected == tx_body.to_primitive()
+
+
+def test_tx_builder_no_output(chain_context):
+    tx_builder = TransactionBuilder(chain_context)
+    sender = "addr_test1vrm9x2zsux7va6w892g38tvchnzahvcd9tykqf3ygnmwtaqyfg52x"
+    sender_address = Address.from_primitive(sender)
+
+    input_amount = 10000000
+
+    tx_in1 = TransactionInput.from_primitive([b"1" * 32, 3])
+    tx_out1 = TransactionOutput.from_primitive([sender, input_amount])
+    utxo1 = UTxO(tx_in1, tx_out1)
+
+    tx_builder.add_input(utxo1)
+
+    tx_body = tx_builder.build(change_address=sender_address, merge_change=True)
+
+    expected = {
+        0: [[b"11111111111111111111111111111111", 3]],
+        1: [
+            [sender_address.to_primitive(), 9836215],
+        ],
+        2: 163785,
+    }
+
+    assert expected == tx_body.to_primitive()
+
+
+def test_tx_builder_merge_change_to_output(chain_context):
+    tx_builder = TransactionBuilder(chain_context)
+    sender = "addr_test1vrm9x2zsux7va6w892g38tvchnzahvcd9tykqf3ygnmwtaqyfg52x"
+    sender_address = Address.from_primitive(sender)
+
+    input_amount = 10000000
+
+    tx_in1 = TransactionInput.from_primitive([b"1" * 32, 3])
+    tx_out1 = TransactionOutput.from_primitive([sender, input_amount])
+    utxo1 = UTxO(tx_in1, tx_out1)
+
+    tx_builder.add_input(utxo1)
+    tx_builder.add_output(TransactionOutput.from_primitive([sender, 10000]))
+
+    tx_body = tx_builder.build(change_address=sender_address, merge_change=True)
+
+    expected = {
+        0: [[b"11111111111111111111111111111111", 3]],
+        1: [
+            [sender_address.to_primitive(), 9836215],
+        ],
+        2: 163785,
+    }
+
+    assert expected == tx_body.to_primitive()
+
+
+def test_tx_builder_merge_change_to_output_2(chain_context):
+    tx_builder = TransactionBuilder(chain_context)
+    sender = "addr_test1vrm9x2zsux7va6w892g38tvchnzahvcd9tykqf3ygnmwtaqyfg52x"
+    sender_address = Address.from_primitive(sender)
+    receiver = "addr_test1vr2p8st5t5cxqglyjky7vk98k7jtfhdpvhl4e97cezuhn0cqcexl7"
+    receiver_address = Address.from_primitive(receiver)
+
+    input_amount = 10000000
+
+    tx_in1 = TransactionInput.from_primitive([b"1" * 32, 3])
+    tx_out1 = TransactionOutput.from_primitive([sender, input_amount])
+    utxo1 = UTxO(tx_in1, tx_out1)
+
+    tx_builder.add_input(utxo1)
+    tx_builder.add_output(TransactionOutput.from_primitive([sender, 10000]))
+    tx_builder.add_output(TransactionOutput.from_primitive([receiver, 10000]))
+    tx_builder.add_output(TransactionOutput.from_primitive([sender, 0]))
+
+    tx_body = tx_builder.build(change_address=sender_address, merge_change=True)
+
+    expected = {
+        0: [[b"11111111111111111111111111111111", 3]],
+        1: [
+            [sender_address.to_primitive(), 10000],
+            [receiver_address.to_primitive(), 10000],
+            [sender_address.to_primitive(), 9813135],
+        ],
+        2: 166865,
+    }
+
+    assert expected == tx_body.to_primitive()
+
+
+def test_tx_builder_merge_change_to_zero_amount_output(chain_context):
+    tx_builder = TransactionBuilder(chain_context)
+    sender = "addr_test1vrm9x2zsux7va6w892g38tvchnzahvcd9tykqf3ygnmwtaqyfg52x"
+    sender_address = Address.from_primitive(sender)
+
+    input_amount = 10000000
+
+    tx_in1 = TransactionInput.from_primitive([b"1" * 32, 3])
+    tx_out1 = TransactionOutput.from_primitive([sender, input_amount])
+    utxo1 = UTxO(tx_in1, tx_out1)
+
+    tx_builder.add_input(utxo1)
+    tx_builder.add_output(TransactionOutput.from_primitive([sender, 0]))
+
+    tx_body = tx_builder.build(change_address=sender_address, merge_change=True)
+
+    expected = {
+        0: [[b"11111111111111111111111111111111", 3]],
+        1: [
+            [sender_address.to_primitive(), 9836215],
+        ],
+        2: 163785,
+    }
+
+    assert expected == tx_body.to_primitive()