also provide methods to sort psbt_{in|out} inplace

btclib/psbt/psbt.py

@@ -18,7 +18,7 @@
 import random
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Any, Callable, Mapping, Sequence, TypeVar
+from typing import Any, Callable, List, Mapping, Sequence, Tuple, TypeVar, cast
 
 from btclib.alias import Octets, String
 from btclib.bip32 import (
@@ -323,6 +323,12 @@ def from_tx(cls: type[Psbt], tx: Tx, check_validity: bool = True) -> Psbt:
             check_validity,
         )
 
+    def sort_inputs(self, key: Callable[[PsbtIn], int] | None = None) -> None:
+        self.inputs, self.tx.vin = _sort_together(self.inputs, self.tx.vin, key=key)
+
+    def sort_outputs(self, key: Callable[[PsbtOut], int] | None = None) -> None:
+        self.outputs, self.tx.vout = _sort_together(self.outputs, self.tx.vout, key=key)
+
 
 def _combine_field(
     psbt_map: PsbtIn | PsbtOut | Psbt, out: PsbtIn | PsbtOut | Psbt, key: str
@@ -465,11 +471,12 @@ def _sort_together(
     if len(list_a) != len(list_b):
         raise ValueError("All lists need to have the same length")
     tmp = list(zip(list_a, list_b))
-    if key:
-        tmp.sort(key=lambda t: key(t[0]))
+    if key is not None:
+        tmp.sort(key=lambda t: cast(Callable[[T], int], key)(t[0]))
     else:
         random.shuffle(tmp)
-    return tuple(zip(*tmp))
+    list_a, list_b = cast(Tuple[List[T], List[Any]], tuple(zip(*tmp)))
+    return (list_a, list_b)
 
 
 def _ensure_consistency(psbts: Sequence[Psbt]) -> None:
@@ -494,10 +501,9 @@ def join_psbts(
     sort_outputs: Callable[[PsbtOut], int] | None = None,
 ) -> Psbt:
     """Join multiple psbts into a single one by merging inputsand outputs
-
     inputs/outputs are shuffled by default. If shuffle=False, they are simply
     concatenated in the same order as psbts are specified. A specific ordering can be
-    specified via sort{inputs|outputs}, which present overwrite shuffle
+    specified via sort_{inputs|outputs}, which overwrite shuffle when present.
     """
     _ensure_consistency(psbts)