@@ -250,24 +250,47 @@ def add_script_input(
250250 self ._consolidate_redeemer (redeemer )
251251 self ._inputs_to_redeemers [utxo ] = redeemer
252252
253+ input_script_hash = utxo .output .address .payment_part
254+ if not isinstance (input_script_hash , ScriptHash ):
255+ raise InvalidArgumentException (
256+ f"Expect the payment part of the address to be of a script (type ScriptHash), "
257+ f"but got { type (input_script_hash )} instead."
258+ )
259+
260+ # collect potential scripts to fulfill the input
261+ candidate_scripts : List [
262+ Tuple [ScriptType , Optional [UTxO ]]
263+ ] = []
253264 if utxo .output .script :
254- self ._inputs_to_scripts [utxo ] = utxo .output .script
255- self .reference_inputs .add (utxo )
256- self ._reference_scripts .append (utxo .output .script )
265+ candidate_scripts .append ((utxo .output .script , utxo ))
257266 elif not script :
258267 for i in self .context .utxos (utxo .output .address ):
259268 if i .output .script :
260- self ._inputs_to_scripts [utxo ] = i .output .script
261- self .reference_inputs .add (i )
262- self ._reference_scripts .append (i .output .script )
263- break
269+ candidate_scripts .append ((i .output .script , i ))
264270 elif isinstance (script , UTxO ):
265- assert script .output .script is not None
266- self ._inputs_to_scripts [utxo ] = script .output .script
267- self .reference_inputs .add (script )
268- self ._reference_scripts .append (script .output .script )
271+ if script .output .script is None :
272+ raise InvalidArgumentException (
273+ "Expect the output of the UTxO to have a script, "
274+ "but got None instead."
275+ )
276+ candidate_scripts .append ((script .output .script , script ))
269277 else :
270- self ._inputs_to_scripts [utxo ] = script
278+ candidate_scripts .append ((script , None ))
279+
280+ found_valid_script = False
281+ for candidate_script , candidate_utxo in candidate_scripts :
282+ if script_hash (candidate_script ) != input_script_hash :
283+ continue
284+ found_valid_script = True
285+ self ._inputs_to_scripts [utxo ] = candidate_script
286+ if candidate_utxo is not None :
287+ self .reference_inputs .add (candidate_utxo )
288+ self ._reference_scripts .append (candidate_script )
289+ if not found_valid_script :
290+ raise InvalidArgumentException (
291+ f"Cannot find a valid script to fulfill the input UTxO: { utxo } ."
292+ "Supplied scripts do not match the payment part of the input address."
293+ )
271294
272295 self .inputs .append (utxo )
273296 return self
0 commit comments