@@ -16,6 +16,11 @@ use tokio::sync::{mpsc, oneshot};
1616
1717use crate :: logger:: Logger ;
1818
19+ type TasksAndContexts = (
20+ HashMap < String , Arc < Py < PyAny > > > ,
21+ HashMap < String , Arc < Py < PyAny > > > ,
22+ ) ;
23+
1924#[ derive( Debug ) ]
2025pub struct TaskRegistry {
2126 tasks : Arc < RwLock < HashMap < String , Arc < Py < PyAny > > > > > ,
@@ -48,56 +53,51 @@ impl TaskRegistry {
4853 let real_module_path = real_module_path. unwrap ( ) ;
4954 let module_dir = project_root. to_string_lossy ( ) . to_string ( ) ;
5055
51- let ( tasks, contexts) = Python :: attach (
52- |py| -> Result < (
53- HashMap < String , Arc < Py < PyAny > > > ,
54- HashMap < String , Arc < Py < PyAny > > > ,
55- ) > {
56- let module = PyModule :: from_code (
57- py,
58- script_cstr. as_c_str ( ) ,
59- filename. as_c_str ( ) ,
60- module_name. as_c_str ( ) ,
61- )
62- . map_err ( |e| anyhow ! ( "Failed to import python module: {}" , e) ) ?;
63-
64- let registry: Bound < ' _ , PyDict > = module
65- . getattr ( "get_registry" )
66- . map_err ( |e| anyhow ! ( "Failed to get 'get_registry' script: {}" , e) ) ?
67- . call1 ( ( real_module_path, queue_name, module_dir) )
68- . map_err ( |e| anyhow ! ( "Failed to get tasks: {}" , e) ) ?
69- . cast_into :: < PyDict > ( )
70- . map_err ( |_| anyhow ! ( "Failed to cast result to a Python Dictionary" ) ) ?;
71-
72- let tasks: HashMap < String , Arc < Py < PyAny > > > = registry
73- . get_item ( "tasks" ) ?
74- . expect ( "tasks missing" )
75- . cast :: < PyDict > ( )
76- . map_err ( |e| anyhow ! ( "tasks is not a dict: {}" , e) ) ?
77- . iter ( )
78- . filter_map ( |( key, value) | {
79- let name: String = key. extract ( ) . ok ( ) ?;
80- let func: Py < PyAny > = value. unbind ( ) ;
81- Some ( ( name, Arc :: new ( func) ) )
82- } )
83- . collect ( ) ;
84-
85- let contexts: HashMap < String , Arc < Py < PyAny > > > = registry
86- . get_item ( "contexts" ) ?
87- . expect ( "contexts missing" )
88- . cast :: < PyDict > ( )
89- . map_err ( |e| anyhow ! ( "contexts is not a dict: {}" , e) ) ?
90- . iter ( )
91- . filter_map ( |( key, value) : ( Bound < PyAny > , Bound < PyAny > ) | {
92- let name: String = key. extract ( ) . ok ( ) ?;
93- let func: Py < PyAny > = value. unbind ( ) ;
94- Some ( ( name, Arc :: new ( func) ) )
95- } )
96- . collect ( ) ;
97-
98- Ok ( ( tasks, contexts) )
99- } ,
100- ) ?;
56+ let ( tasks, contexts) = Python :: attach ( |py| -> Result < TasksAndContexts > {
57+ let module = PyModule :: from_code (
58+ py,
59+ script_cstr. as_c_str ( ) ,
60+ filename. as_c_str ( ) ,
61+ module_name. as_c_str ( ) ,
62+ )
63+ . map_err ( |e| anyhow ! ( "Failed to import python module: {}" , e) ) ?;
64+
65+ let registry: Bound < ' _ , PyDict > = module
66+ . getattr ( "get_registry" )
67+ . map_err ( |e| anyhow ! ( "Failed to get 'get_registry' script: {}" , e) ) ?
68+ . call1 ( ( real_module_path, queue_name, module_dir) )
69+ . map_err ( |e| anyhow ! ( "Failed to get tasks: {}" , e) ) ?
70+ . cast_into :: < PyDict > ( )
71+ . map_err ( |_| anyhow ! ( "Failed to cast result to a Python Dictionary" ) ) ?;
72+
73+ let tasks: HashMap < String , Arc < Py < PyAny > > > = registry
74+ . get_item ( "tasks" ) ?
75+ . expect ( "tasks missing" )
76+ . cast :: < PyDict > ( )
77+ . map_err ( |e| anyhow ! ( "tasks is not a dict: {}" , e) ) ?
78+ . iter ( )
79+ . filter_map ( |( key, value) | {
80+ let name: String = key. extract ( ) . ok ( ) ?;
81+ let func: Py < PyAny > = value. unbind ( ) ;
82+ Some ( ( name, Arc :: new ( func) ) )
83+ } )
84+ . collect ( ) ;
85+
86+ let contexts: HashMap < String , Arc < Py < PyAny > > > = registry
87+ . get_item ( "contexts" ) ?
88+ . expect ( "contexts missing" )
89+ . cast :: < PyDict > ( )
90+ . map_err ( |e| anyhow ! ( "contexts is not a dict: {}" , e) ) ?
91+ . iter ( )
92+ . filter_map ( |( key, value) : ( Bound < PyAny > , Bound < PyAny > ) | {
93+ let name: String = key. extract ( ) . ok ( ) ?;
94+ let func: Py < PyAny > = value. unbind ( ) ;
95+ Some ( ( name, Arc :: new ( func) ) )
96+ } )
97+ . collect ( ) ;
98+
99+ Ok ( ( tasks, contexts) )
100+ } ) ?;
101101
102102 Ok ( Self {
103103 tasks : Arc :: new ( RwLock :: new ( tasks) ) ,
0 commit comments