Skip to content

Commit aaa8658

Browse files
committed
Add _asyncio
1 parent 2e0b765 commit aaa8658

2 files changed

Lines changed: 131 additions & 0 deletions

File tree

crates/stdlib/src/_asyncio.rs

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
//! _asyncio module - provides native asyncio support
2+
//!
3+
//! This module uses lazy attribute loading to avoid circular import issues.
4+
//! When attributes like Future or Task are accessed, they are imported from
5+
//! the Python asyncio modules on demand.
6+
7+
use crate::vm::{
8+
Context, Py, PyResult, VirtualMachine,
9+
builtins::{PyModule, PyModuleDef, PyModuleSlots},
10+
common::static_cell,
11+
compiler,
12+
};
13+
14+
pub(crate) const MODULE_NAME: &str = "_asyncio";
15+
16+
static_cell! {
17+
static DEF: PyModuleDef;
18+
}
19+
20+
/// Returns the module definition for multi-phase initialization.
21+
pub(crate) fn module_def(ctx: &Context) -> &'static PyModuleDef {
22+
DEF.get_or_init(|| PyModuleDef {
23+
name: ctx.intern_str(MODULE_NAME),
24+
doc: Some(ctx.intern_str("Accelerator module for asyncio")),
25+
methods: &[],
26+
slots: PyModuleSlots {
27+
create: None,
28+
exec: Some(exec_module),
29+
},
30+
})
31+
}
32+
33+
/// Exec phase: Set up the module __getattr__ for lazy attribute loading.
34+
/// We don't import asyncio here to avoid circular import issues.
35+
fn exec_module(vm: &VirtualMachine, module: &Py<PyModule>) -> PyResult<()> {
36+
// Define __getattr__ function that lazily loads attributes from asyncio
37+
let getattr_code = r#"
38+
def __getattr__(name):
39+
import sys
40+
_asyncio = sys.modules['_asyncio']
41+
42+
# Mapping of _asyncio attributes to their sources
43+
if name == 'Future':
44+
from asyncio.futures import _PyFuture
45+
setattr(_asyncio, 'Future', _PyFuture)
46+
return _PyFuture
47+
elif name == 'Task':
48+
from asyncio.tasks import _PyTask
49+
setattr(_asyncio, 'Task', _PyTask)
50+
return _PyTask
51+
elif name == 'current_task':
52+
from asyncio.tasks import _py_current_task
53+
setattr(_asyncio, 'current_task', _py_current_task)
54+
return _py_current_task
55+
elif name == '_register_task':
56+
from asyncio.tasks import _py_register_task
57+
setattr(_asyncio, '_register_task', _py_register_task)
58+
return _py_register_task
59+
elif name == '_register_eager_task':
60+
from asyncio.tasks import _py_register_eager_task
61+
setattr(_asyncio, '_register_eager_task', _py_register_eager_task)
62+
return _py_register_eager_task
63+
elif name == '_unregister_task':
64+
from asyncio.tasks import _py_unregister_task
65+
setattr(_asyncio, '_unregister_task', _py_unregister_task)
66+
return _py_unregister_task
67+
elif name == '_unregister_eager_task':
68+
from asyncio.tasks import _py_unregister_eager_task
69+
setattr(_asyncio, '_unregister_eager_task', _py_unregister_eager_task)
70+
return _py_unregister_eager_task
71+
elif name == '_enter_task':
72+
from asyncio.tasks import _py_enter_task
73+
setattr(_asyncio, '_enter_task', _py_enter_task)
74+
return _py_enter_task
75+
elif name == '_leave_task':
76+
from asyncio.tasks import _py_leave_task
77+
setattr(_asyncio, '_leave_task', _py_leave_task)
78+
return _py_leave_task
79+
elif name == '_swap_current_task':
80+
from asyncio.tasks import _py_swap_current_task
81+
setattr(_asyncio, '_swap_current_task', _py_swap_current_task)
82+
return _py_swap_current_task
83+
elif name == '_scheduled_tasks':
84+
from asyncio.tasks import _scheduled_tasks
85+
setattr(_asyncio, '_scheduled_tasks', _scheduled_tasks)
86+
return _scheduled_tasks
87+
elif name == '_eager_tasks':
88+
from asyncio.tasks import _eager_tasks
89+
setattr(_asyncio, '_eager_tasks', _eager_tasks)
90+
return _eager_tasks
91+
elif name == '_current_tasks':
92+
from asyncio.tasks import _current_tasks
93+
setattr(_asyncio, '_current_tasks', _current_tasks)
94+
return _current_tasks
95+
elif name == '_get_running_loop':
96+
from asyncio.events import _py__get_running_loop
97+
setattr(_asyncio, '_get_running_loop', _py__get_running_loop)
98+
return _py__get_running_loop
99+
elif name == '_set_running_loop':
100+
from asyncio.events import _py__set_running_loop
101+
setattr(_asyncio, '_set_running_loop', _py__set_running_loop)
102+
return _py__set_running_loop
103+
elif name == 'get_running_loop':
104+
from asyncio.events import _py_get_running_loop
105+
setattr(_asyncio, 'get_running_loop', _py_get_running_loop)
106+
return _py_get_running_loop
107+
elif name == 'get_event_loop':
108+
from asyncio.events import _py_get_event_loop
109+
setattr(_asyncio, 'get_event_loop', _py_get_event_loop)
110+
return _py_get_event_loop
111+
112+
raise AttributeError(f"module '_asyncio' has no attribute '{name}'")
113+
"#;
114+
115+
// Execute the code in the module's namespace
116+
let code = vm
117+
.compile(getattr_code, compiler::Mode::Exec, "<_asyncio>".to_owned())
118+
.map_err(|e| vm.new_syntax_error(&e, Some(getattr_code)))?;
119+
120+
let scope = vm.new_scope_with_builtins();
121+
vm.run_code_obj(code, scope.clone())?;
122+
123+
// Get the __getattr__ function and set it on the module
124+
if let Ok(getattr_func) = scope.globals.get_item("__getattr__", vm) {
125+
module.set_attr("__getattr__", getattr_func, vm)?;
126+
}
127+
128+
Ok(())
129+
}

crates/stdlib/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ extern crate alloc;
1111
#[macro_use]
1212
pub(crate) mod macros;
1313

14+
mod _asyncio;
1415
pub mod array;
1516
mod binascii;
1617
mod bisect;
@@ -112,6 +113,7 @@ use crate::vm::{Context, builtins};
112113
/// allowing safe circular imports.
113114
pub fn stdlib_module_defs(ctx: &Context) -> Vec<&'static builtins::PyModuleDef> {
114115
vec![
116+
_asyncio::module_def(ctx),
115117
_opcode::module_def(ctx),
116118
array::module_def(ctx),
117119
binascii::module_def(ctx),

0 commit comments

Comments
 (0)