@@ -3,7 +3,7 @@ mod receipt;
33mod session;
44
55use crate :: image:: Image ;
6- use crate :: receipt:: Receipt ;
6+ use crate :: receipt:: { Receipt , ExitStatus , ExitKind } ;
77use crate :: session:: { ExitCode , SessionInfo } ;
88use pyo3:: prelude:: * ;
99use risc0_zkvm:: { default_prover, ExecutorEnv , ProverOpts } ;
@@ -87,6 +87,22 @@ fn prove_with_opts(_py: Python<'_>, image: &Image, input_bytes: &Bound<'_, PyAny
8787// Advanced functions removed - segments are no longer exposed
8888// If needed in future, these could work with Receipt types instead
8989
90+ /// Compute the expected image ID from an ELF file as hex string
91+ ///
92+ /// Args:
93+ /// elf_bytes: The ELF binary to compute ID from
94+ ///
95+ /// Returns:
96+ /// 64-character hex string of the image ID
97+ #[ pyfunction]
98+ fn compute_image_id_hex ( elf_bytes : Vec < u8 > ) -> PyResult < String > {
99+ let image_id = risc0_binfmt:: compute_image_id ( & elf_bytes)
100+ . map_err ( |e| PyErr :: new :: < pyo3:: exceptions:: PyValueError , _ > (
101+ format ! ( "Failed to compute image ID: {}" , e)
102+ ) ) ?;
103+ Ok ( hex:: encode ( image_id) )
104+ }
105+
90106
91107
92108
@@ -96,11 +112,14 @@ fn _rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
96112 m. add_class :: < ExitCode > ( ) ?;
97113 m. add_class :: < SessionInfo > ( ) ?;
98114 m. add_class :: < Receipt > ( ) ?;
115+ m. add_class :: < ExitStatus > ( ) ?;
116+ m. add_class :: < ExitKind > ( ) ?;
99117
100118 // Core API functions
101119 m. add_function ( wrap_pyfunction ! ( load_image, m) ?) ?;
102120 m. add_function ( wrap_pyfunction ! ( prove, m) ?) ?;
103121 m. add_function ( wrap_pyfunction ! ( prove_with_opts, m) ?) ?;
122+ m. add_function ( wrap_pyfunction ! ( compute_image_id_hex, m) ?) ?;
104123
105124 // Optional debugging function
106125 m. add_function ( wrap_pyfunction ! ( dry_run, m) ?) ?;
0 commit comments