Skip to content

Commit 9850b5f

Browse files
coolreader18coolreader18
committed
Implement _thread._local
1 parent e14adc4 commit 9850b5f

3 files changed

Lines changed: 99 additions & 5 deletions

File tree

‎vm/src/obj/objmodule.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,14 +67,15 @@ impl PyModuleRef {
6767
vm.generic_getattribute_opt(
6868
self.as_object().clone(),
6969
PyString::from("__name__").into_ref(vm),
70+
None,
7071
)
7172
.unwrap_or(None)
7273
.and_then(|obj| obj.payload::<PyString>().map(|s| s.as_str().to_owned()))
7374
}
7475

7576
#[pymethod(magic)]
7677
fn getattribute(self, name: PyStringRef, vm: &VirtualMachine) -> PyResult {
77-
vm.generic_getattribute_opt(self.as_object().clone(), name.clone())?
78+
vm.generic_getattribute_opt(self.as_object().clone(), name.clone(), None)?
7879
.ok_or_else(|| {
7980
let module_name = if let Some(name) = self.name(vm) {
8081
format!(" '{}'", name)

‎vm/src/stdlib/thread.rs‎

Lines changed: 92 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,22 @@
22
use crate::exceptions;
33
use crate::function::{Args, KwArgs, OptionalArg, PyFuncArgs};
44
use crate::obj::objdict::PyDictRef;
5+
use crate::obj::objstr::PyStringRef;
56
use crate::obj::objtuple::PyTupleRef;
67
use crate::obj::objtype::PyClassRef;
78
use crate::pyobject::{
8-
Either, IdProtocol, PyCallable, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue,
9-
TypeProtocol,
9+
Either, IdProtocol, ItemProtocol, PyCallable, PyClassImpl, PyObjectRef, PyRef, PyResult,
10+
PyValue, TypeProtocol,
1011
};
1112
use crate::vm::VirtualMachine;
1213

14+
use crossbeam_utils::atomic::AtomicCell;
1315
use parking_lot::{
1416
lock_api::{RawMutex as RawMutexT, RawMutexTimed, RawReentrantMutex},
1517
RawMutex, RawThreadId,
1618
};
1719
use std::cell::RefCell;
20+
use std::collections::HashMap;
1821
use std::io::Write;
1922
use std::time::Duration;
2023
use std::{fmt, thread};
@@ -264,12 +267,99 @@ fn thread_count(vm: &VirtualMachine) -> usize {
264267
vm.state.thread_count.load()
265268
}
266269

270+
static LOCAL_DATA_KEY: AtomicCell<usize> = AtomicCell::new(0);
271+
thread_local!(static LOCAL_DATA: RefCell<HashMap<usize, PyDictRef>> = RefCell::default());
272+
273+
#[pyclass(name = "_local")]
274+
#[derive(Debug)]
275+
struct PyLocal {
276+
key: usize,
277+
}
278+
279+
impl PyValue for PyLocal {
280+
fn class(vm: &VirtualMachine) -> PyClassRef {
281+
vm.class("_thread", "_local")
282+
}
283+
}
284+
285+
impl Drop for PyLocal {
286+
fn drop(&mut self) {
287+
LOCAL_DATA.with(|map| map.borrow_mut().remove(&self.key));
288+
}
289+
}
290+
291+
#[pyimpl(flags(BASETYPE))]
292+
impl PyLocal {
293+
fn ldict(&self, vm: &VirtualMachine) -> PyDictRef {
294+
LOCAL_DATA.with(|map| {
295+
map.borrow_mut()
296+
.entry(self.key)
297+
.or_insert_with(|| vm.ctx.new_dict())
298+
.clone()
299+
})
300+
}
301+
302+
#[pyslot]
303+
fn tp_new(cls: PyClassRef, _args: PyFuncArgs, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
304+
PyLocal {
305+
key: LOCAL_DATA_KEY.fetch_add(1),
306+
}
307+
.into_ref_with_type(vm, cls)
308+
}
309+
310+
#[pymethod(magic)]
311+
fn getattribute(zelf: PyRef<Self>, attr: PyStringRef, vm: &VirtualMachine) -> PyResult {
312+
let ldict = zelf.ldict(vm);
313+
if attr.as_str() == "__dict__" {
314+
Ok(ldict.into_object())
315+
} else {
316+
let zelf = zelf.into_object();
317+
vm.generic_getattribute_opt(zelf.clone(), attr.clone(), Some(ldict))?
318+
.ok_or_else(|| {
319+
vm.new_attribute_error(format!("{} has no attribute '{}'", zelf, attr))
320+
})
321+
}
322+
}
323+
324+
#[pymethod(magic)]
325+
fn setattr(
326+
zelf: PyRef<Self>,
327+
attr: PyStringRef,
328+
value: PyObjectRef,
329+
vm: &VirtualMachine,
330+
) -> PyResult<()> {
331+
if attr.as_str() == "__dict__" {
332+
Err(vm.new_attribute_error(format!(
333+
"{} attribute '__dict__' is read-only",
334+
zelf.as_object()
335+
)))
336+
} else {
337+
zelf.ldict(vm).set_item(attr.as_object(), value, vm)?;
338+
Ok(())
339+
}
340+
}
341+
342+
#[pymethod(magic)]
343+
fn delattr(zelf: PyRef<Self>, attr: PyStringRef, vm: &VirtualMachine) -> PyResult<()> {
344+
if attr.as_str() == "__dict__" {
345+
Err(vm.new_attribute_error(format!(
346+
"{} attribute '__dict__' is read-only",
347+
zelf.as_object()
348+
)))
349+
} else {
350+
zelf.ldict(vm).del_item(attr.as_object(), vm)?;
351+
Ok(())
352+
}
353+
}
354+
}
355+
267356
pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
268357
let ctx = &vm.ctx;
269358

270359
py_module!(vm, "_thread", {
271360
"RLock" => PyRLock::make_class(ctx),
272361
"LockType" => PyLock::make_class(ctx),
362+
"_local" => PyLocal::make_class(ctx),
273363
"get_ident" => ctx.new_function(thread_get_ident),
274364
"allocate_lock" => ctx.new_function(thread_allocate_lock),
275365
"start_new_thread" => ctx.new_function(thread_start_new_thread),

‎vm/src/vm.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1001,7 +1001,7 @@ impl VirtualMachine {
10011001
}
10021002

10031003
pub fn generic_getattribute(&self, obj: PyObjectRef, name: PyStringRef) -> PyResult {
1004-
self.generic_getattribute_opt(obj.clone(), name.clone())?
1004+
self.generic_getattribute_opt(obj.clone(), name.clone(), None)?
10051005
.ok_or_else(|| self.new_attribute_error(format!("{} has no attribute '{}'", obj, name)))
10061006
}
10071007

@@ -1010,6 +1010,7 @@ impl VirtualMachine {
10101010
&self,
10111011
obj: PyObjectRef,
10121012
name_str: PyStringRef,
1013+
dict: Option<PyDictRef>,
10131014
) -> PyResult<Option<PyObjectRef>> {
10141015
let name = name_str.as_str();
10151016
let cls = obj.class();
@@ -1023,7 +1024,9 @@ impl VirtualMachine {
10231024
}
10241025
}
10251026

1026-
let attr = if let Some(dict) = obj.dict() {
1027+
let dict = dict.or_else(|| obj.dict());
1028+
1029+
let attr = if let Some(dict) = dict {
10271030
dict.get_item_option(name_str.as_str(), self)?
10281031
} else {
10291032
None

0 commit comments

Comments
 (0)