|
| 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 | +} |
0 commit comments