Repository: RustPython/RustPython (language: rust)

Issue:
[RFC] Implement PEP 567 (`contextvars`)
## Summary

<!-- A quick overview of your suggestion -->
The `contextvars` module, from PEP 567, is necessary to upgrade `asyncio` (and `decimal`) to CPython 3.8. I started #2950 to implement `contextvars`, but the PR has since stalled. (@DimitrisJim urged me to open this issue.)

## Detailed Explanation and Rationale

<!-- Elaborate on your suggestion in all its details -->
[PEP 567](https://www.python.org/dev/peps/pep-0567/), implemented in [CPython 3.7](https://docs.python.org/3/whatsnew/3.7.html#contextvars), introduced the `contextvars` module and context variables, which are conceptually similar to thread-local variables. However, unlike thread-local variables, context variables correctly support asynchronous code. Following that implementation, the `asyncio` and `decimal` modules were re-written to take advantage of this new capability. As part of our endeavors to get as much of the codebase covered by the CPython test suite as possible (see #1671), and out of a desire support CPython 3.8 or later (see the README), I think it's necessary to add the `test_asyncio` module to the CPython test suite. The addition of the CPython 3.8 version of `test_asyncio` would benefit from upgrading the `asyncio` module itself to that same version. (@coolreader18 and I have noticed that test failures were often alleviated by upgrading the relevant Python module.) However, `asyncio` has a dependency — `contextvars` — that we don't have yet. Accordingly, I have started a draft PR (#2950) to implement this module here, but I got stymied when I was drafting the Rust implementation because I am unfamiliar with how RustPython deals with thread safety. For context, here are the parts of the C implementation that I think are relevant:
- https://github.com/python/cpython/blob/3.8/Include/context.h
- https://github.com/python/cpython/blob/3.8/Python/context.c
- https://github.com/python/cpython/blob/3.8/Python/clinic/context.c.h
- https://github.com/python/cpython/blob/3.8/Modules/_contextvarsmodule.c
- https://github.com/python/cpython/blob/3.8/Modules/clinic/_contextvarsmodule.c.h
- https://github.com/python/cpython/blob/3.8/Include/internal/pycore_context.h


## Unresolved Questions

<!-- What would you like feedback on for fleshing out your suggestion? -->
- What is the meaning of `PyContext` (defined in [`pyobject.rs`](https://github.com/RustPython/RustPython/blob/main/vm/src/pyobject.rs)), and does it have any bearing on the scope of PEP 567?
- The scaffolding is present, but how should it be fleshed out?
- How would each `Context` object relate to a thread in RustPython?



Discussion / hints:

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index a7d72388f6..11c93a636e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2241,6 +2241,7 @@ dependencies = [
  "foreign-types-shared",
  "gethostname",
  "hex",
+ "indexmap 2.2.6",
  "itertools 0.11.0",
  "junction",
  "libc",
@@ -2277,6 +2278,7 @@ dependencies = [
  "socket2",
  "system-configuration",
  "termios",
+ "thread_local",
  "ucd",
  "unic-char-property",
  "unic-normal",
diff --git a/common/src/hash.rs b/common/src/hash.rs
index 558b0fe15f..4e87eff799 100644
--- a/common/src/hash.rs
+++ b/common/src/hash.rs
@@ -88,6 +88,13 @@ impl HashSecret {
     }
 }
 
+#[inline]
+pub fn hash_pointer(value: usize) -> PyHash {
+    // TODO: 32bit?
+    let hash = (value >> 4) | value;
+    hash as _
+}
+
 #[inline]
 pub fn hash_float(value: f64) -> Option<PyHash> {
     // cpython _Py_HashDouble
diff --git a/stdlib/Cargo.toml b/stdlib/Cargo.toml
index f919efc766..9a6be30987 100644
--- a/stdlib/Cargo.toml
+++ b/stdlib/Cargo.toml
@@ -28,6 +28,7 @@ cfg-if = { workspace = true }
 crossbeam-utils = { workspace = true }
 hex = { workspace = true }
 itertools = { workspace = true }
+indexmap = { workspace = true }
 libc = { workspace = true }
 nix = { workspace = true }
 num-complex = { workspace = true }
@@ -37,6 +38,7 @@ num-traits = { workspace = true }
 num_enum = { workspace = true }
 once_cell = { workspace = true }
 parking_lot = { workspace = true }
+thread_local = { workspace = true }
 
 memchr = { workspace = true }
 base64 = "0.13.0"
diff --git a/stdlib/src/contextvars.rs b/stdlib/src/contextvars.rs
index d808d1d08e..8e1bfca005 100644
--- a/stdlib/src/contextvars.rs
+++ b/stdlib/src/contextvars.rs
@@ -1,49 +1,217 @@
-pub(crate) use _contextvars::make_module;
+use crate::vm::{builtins::PyModule, class::StaticType, PyRef, VirtualMachine};
+use _contextvars::PyContext;
+use std::cell::RefCell;
+
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
+    let module = _contextvars::make_module(vm);
+    let token_type = module.get_attr("Token", vm).unwrap();
+    token_type
+        .set_attr(
+            "MISSING",
+            _contextvars::ContextTokenMissing::static_type().to_owned(),
+            vm,
+        )
+        .unwrap();
+    module
+}
+
+thread_local! {
+    // TODO: Vec doesn't seem to match copy behavior
+    static CONTEXTS: RefCell<Vec<PyRef<PyContext>>> = RefCell::default();
+}
 
 #[pymodule]
 mod _contextvars {
     use crate::vm::{
-        builtins::{PyFunction, PyStrRef, PyTypeRef},
+        atomic_func,
+        builtins::{PyStrRef, PyTypeRef},
+        class::StaticType,
+        common::hash::PyHash,
         function::{ArgCallable, FuncArgs, OptionalArg},
-        types::{Initializer, Representable},
-        Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
+        protocol::{PyMappingMethods, PySequenceMethods},
+        types::{AsMapping, AsSequence, Constructor, Hashable, Representable},
+        AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
+    };
+    use crossbeam_utils::atomic::AtomicCell;
+    use indexmap::IndexMap;
+    use once_cell::sync::Lazy;
+    use std::{
+        cell::{Cell, RefCell, UnsafeCell},
+        sync::atomic::Ordering,
     };
 
+    // TODO: Real hamt implementation
+    type Hamt = IndexMap<PyRef<ContextVar>, PyObjectRef, ahash::RandomState>;
+
+    #[pyclass(no_attr, name = "Hamt", module = "contextvars")]
+    #[derive(Debug, PyPayload)]
+    pub(crate) struct HamtObject {
+        hamt: RefCell<Hamt>,
+    }
+
+    #[pyclass]
+    impl HamtObject {}
+
+    impl Default for HamtObject {
+        fn default() -> Self {
+            Self {
+                hamt: RefCell::new(Hamt::default()),
+            }
+        }
+    }
+
+    unsafe impl Sync for HamtObject {}
+
+    #[derive(Debug)]
+    struct ContextInner {
+        idx: Cell<usize>,
+        vars: PyRef<HamtObject>,
+        // PyObject *ctx_weakreflist;
+        entered: Cell<bool>,
+    }
+
+    unsafe impl Sync for ContextInner {}
+
     #[pyattr]
     #[pyclass(name = "Context")]
     #[derive(Debug, PyPayload)]
-    struct PyContext {} // not to confuse with vm::Context
+    pub(crate) struct PyContext {
+        // not to confuse with vm::Context
+        inner: ContextInner,
+    }
 
-    #[pyclass(with(Initializer))]
     impl PyContext {
-        #[pymethod]
-        fn run(
-            &self,
-            _callable: ArgCallable,
-            _args: FuncArgs,
-            _vm: &VirtualMachine,
-        ) -> PyResult<PyFunction> {
-            unimplemented!("Context.run is currently under construction")
+        fn empty(vm: &VirtualMachine) -> Self {
+            Self {
+                inner: ContextInner {
+                    idx: Cell::new(usize::MAX),
+                    vars: HamtObject::default().into_ref(&vm.ctx),
+                    entered: Cell::new(false),
+                },
+            }
+        }
+
+        fn borrow_vars(&self) -> impl std::ops::Deref<Target = Hamt> + '_ {
+            self.inner.vars.hamt.borrow()
+        }
+
+        fn borrow_vars_mut(&self) -> impl std::ops::DerefMut<Target = Hamt> + '_ {
+            self.inner.vars.hamt.borrow_mut()
+        }
+
+        fn enter(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
+            if zelf.inner.entered.get() {
+                let msg = format!(
+                    "cannot enter context: {} is already entered",
+                    zelf.as_object().repr(vm)?
+                );
+                return Err(vm.new_runtime_error(msg));
+            }
+
+            super::CONTEXTS.with(|ctxs| {
+                let mut ctxs = ctxs.borrow_mut();
+                zelf.inner.idx.set(ctxs.len());
+                ctxs.push(zelf.to_owned());
+            });
+            zelf.inner.entered.set(true);
+
+            Ok(())
+        }
+
+        fn exit(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
+            if !zelf.inner.entered.get() {
+                let msg = format!(
+                    "cannot exit context: {} is not entered",
+                    zelf.as_object().repr(vm)?
+                );
+                return Err(vm.new_runtime_error(msg));
+            }
+
+            super::CONTEXTS.with(|ctxs| {
+                let mut ctxs = ctxs.borrow_mut();
+                if !ctxs
+                    .last()
+                    .map_or(false, |ctx| ctx.get_id() == zelf.get_id())
+                {
+                    let msg =
+                        "cannot exit context: thread state references a different context object"
+                            .to_owned();
+                    return Err(vm.new_runtime_error(msg));
+                }
+
+                let _ = ctxs.pop();
+                Ok(())
+            })?;
+            zelf.inner.entered.set(false);
+
+            Ok(())
         }
 
+        fn current(vm: &VirtualMachine) -> PyRef<Self> {
+            super::CONTEXTS.with(|ctxs| {
+                let mut ctxs = ctxs.borrow_mut();
+                if let Some(ctx) = ctxs.last() {
+                    ctx.clone()
+                } else {
+                    let ctx = PyContext::empty(vm);
+                    ctx.inner.idx.set(0);
+                    ctx.inner.entered.set(true);
+                    let ctx = ctx.into_ref(&vm.ctx);
+                    ctxs.push(ctx);
+                    ctxs[0].clone()
+                }
+            })
+        }
+
+        fn contains(&self, needle: &Py<ContextVar>) -> PyResult<bool> {
+            let vars = self.borrow_vars();
+            Ok(vars.get(needle).is_some())
+        }
+
+        fn get_inner(&self, needle: &Py<ContextVar>) -> Option<PyObjectRef> {
+            let vars = self.borrow_vars();
+            vars.get(needle).map(|o| o.to_owned())
+        }
+    }
+
+    #[pyclass(with(Constructor, AsMapping, AsSequence))]
+    impl PyContext {
         #[pymethod]
-        fn copy(&self, _vm: &VirtualMachine) -> PyResult<Self> {
-            unimplemented!("Context.copy is currently under construction")
+        fn run(
+            zelf: &Py<Self>,
+            callable: ArgCallable,
+            args: FuncArgs,
+            vm: &VirtualMachine,
+        ) -> PyResult {
+            Self::enter(zelf, vm)?;
+            let result = callable.invoke(args, vm);
+            Self::exit(zelf, vm)?;
+            result
         }
 
-        #[pymethod(magic)]
-        fn getitem(&self, _var: PyObjectRef) -> PyResult<PyObjectRef> {
-            unimplemented!("Context.__getitem__ is currently under construction")
+        #[pymethod]
+        fn copy(&self) -> Self {
+            Self {
+                inner: ContextInner {
+                    idx: Cell::new(usize::MAX),
+                    vars: self.inner.vars.clone(),
+                    entered: Cell::new(false),
+                },
+            }
         }
 
         #[pymethod(magic)]
-        fn contains(&self, _var: PyObjectRef) -> PyResult<bool> {
-            unimplemented!("Context.__contains__ is currently under construction")
+        fn getitem(&self, var: PyRef<ContextVar>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
+            let vars = self.borrow_vars();
+            let item = vars
+                .get(&*var)
+                .ok_or_else(|| vm.new_key_error(var.into()))?;
+            Ok(item.to_owned())
         }
 
         #[pymethod(magic)]
         fn len(&self) -> usize {
-            unimplemented!("Context.__len__ is currently under construction")
+            self.borrow_vars().len()
         }
 
         #[pymethod(magic)]
@@ -54,53 +222,151 @@ mod _contextvars {
         #[pymethod]
         fn get(
             &self,
-            _key: PyObjectRef,
-            _default: OptionalArg<PyObjectRef>,
-        ) -> PyResult<PyObjectRef> {
-            unimplemented!("Context.get is currently under construction")
+            key: PyRef<ContextVar>,
+            default: OptionalArg<PyObjectRef>,
+        ) -> PyResult<Option<PyObjectRef>> {
+            let found = self.get_inner(&key);
+            let result = if let Some(found) = found {
+                Some(found.to_owned())
+            } else {
+                default.into_option()
+            };
+            Ok(result)
         }
 
+        // TODO: wrong return type
         #[pymethod]
-        fn keys(_zelf: PyRef<Self>, _vm: &VirtualMachine) -> Vec<PyObjectRef> {
-            unimplemented!("Context.keys is currently under construction")
+        fn keys(zelf: &Py<Self>) -> Vec<PyObjectRef> {
+            let vars = zelf.borrow_vars();
+            vars.keys().map(|key| key.to_owned().into()).collect()
         }
 
+        // TODO: wrong return type
         #[pymethod]
-        fn values(_zelf: PyRef<Self>, _vm: &VirtualMachine) -> Vec<PyObjectRef> {
-            unimplemented!("Context.values is currently under construction")
+        fn values(zelf: PyRef<Self>) -> Vec<PyObjectRef> {
+            let vars = zelf.borrow_vars();
+            vars.values().map(|value| value.to_owned()).collect()
         }
     }
 
-    impl Initializer for PyContext {
-        type Args = FuncArgs;
+    impl Constructor for PyContext {
+        type Args = ();
+        fn py_new(_cls: PyTypeRef, _args: Self::Args, vm: &VirtualMachine) -> PyResult {
+            Ok(PyContext::empty(vm).into_pyobject(vm))
+        }
+    }
+
+    impl AsMapping for PyContext {
+        fn as_mapping() -> &'static PyMappingMethods {
+            static AS_MAPPING: PyMappingMethods = PyMappingMethods {
+                length: atomic_func!(|mapping, _vm| Ok(PyContext::mapping_downcast(mapping).len())),
+                subscript: atomic_func!(|mapping, needle, vm| {
+                    let needle = needle.try_to_value(vm)?;
+                    let found = PyContext::mapping_downcast(mapping).get_inner(needle);
+                    if let Some(found) = found {
+                        Ok(found.to_owned())
+                    } else {
+                        Err(vm.new_key_error(needle.to_owned().into()))
+                    }
+                }),
+                ass_subscript: AtomicCell::new(None),
+            };
+            &AS_MAPPING
+        }
+    }
 
-        fn init(_obj: PyRef<Self>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<()> {
-            unimplemented!("Context.__init__ is currently under construction")
+    impl AsSequence for PyContext {
+        fn as_sequence() -> &'static PySequenceMethods {
+            static AS_SEQUENCE: Lazy<PySequenceMethods> = Lazy::new(|| PySequenceMethods {
+                contains: atomic_func!(|seq, target, vm| {
+                    let target = target.try_to_value(vm)?;
+                    PyContext::sequence_downcast(seq).contains(target)
+                }),
+                ..PySequenceMethods::NOT_IMPLEMENTED
+            });
+            &AS_SEQUENCE
         }
     }
 
     #[pyattr]
     #[pyclass(name, traverse)]
-    #[derive(Debug, PyPayload)]
+    #[derive(PyPayload)]
     struct ContextVar {
         #[pytraverse(skip)]
-        #[allow(dead_code)] // TODO: RUSTPYTHON
         name: String,
-        #[allow(dead_code)] // TODO: RUSTPYTHON
         default: Option<PyObjectRef>,
+        #[pytraverse(skip)]
+        cached: AtomicCell<Option<ContextVarCache>>,
+        #[pytraverse(skip)]
+        cached_id: std::sync::atomic::AtomicUsize, // cached_tsid in CPython
+        #[pytraverse(skip)]
+        hash: UnsafeCell<PyHash>,
     }
 
-    #[derive(FromArgs)]
-    struct ContextVarOptions {
-        #[pyarg(positional)]
-        #[allow(dead_code)] // TODO: RUSTPYTHON
-        name: PyStrRef,
-        #[pyarg(any, optional)]
-        #[allow(dead_code)] // TODO: RUSTPYTHON
-        default: OptionalArg<PyObjectRef>,
+    impl std::fmt::Debug for ContextVar {
+        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+            f.debug_struct("ContextVar").finish()
+        }
     }
 
-    #[pyclass(with(Initializer, Representable))]
+    unsafe impl Sync for ContextVar {}
+
+    impl PartialEq for ContextVar {
+        fn eq(&self, other: &Self) -> bool {
+            std::ptr::eq(self, other)
+        }
+    }
+    impl Eq for ContextVar {}
+
+    #[derive(Debug)]
+    struct ContextVarCache {
+        object: PyObjectRef, // value; cached in CPython
+        idx: usize,          // Context index; cached_tsver in CPython
+    }
+
+    impl ContextVar {
+        fn delete(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
+            zelf.cached.store(None);
+
+            let ctx = PyContext::current(vm);
+
+            let mut vars = ctx.borrow_vars_mut();
+            if vars.swap_remove(zelf).is_none() {
+                // TODO:
+                // PyErr_SetObject(PyExc_LookupError, (PyObject *)var);
+                let msg = zelf.as_object().repr(vm)?.as_str().to_owned();
+                return Err(vm.new_lookup_error(msg));
+            }
+
+            Ok(())
+        }
+
+        // contextvar_set in CPython
+        fn set_inner(zelf: &Py<Self>, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+            let ctx = PyContext::current(vm);
+
+            let mut vars = ctx.borrow_vars_mut();
+            vars.insert(zelf.to_owned(), value.clone());
+
+            zelf.cached_id.store(ctx.get_id(), Ordering::SeqCst);
+
+            let cache = ContextVarCache {
+                object: value,
+                idx: ctx.inner.idx.get(),
+            };
+            zelf.cached.store(Some(cache));
+
+            Ok(())
+        }
+
+        fn generate_hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyHash {
+            let name_hash = vm.state.hash_secret.hash_str(&zelf.name);
+            let pointer_hash = crate::common::hash::hash_pointer(zelf.as_object().get_id());
+            pointer_hash ^ name_hash
+        }
+    }
+
+    #[pyclass(with(Constructor, Hashable, Representable))]
     impl ContextVar {
         #[pygetset]
         fn name(&self) -> String {
@@ -109,25 +375,102 @@ mod _contextvars {
 
         #[pymethod]
         fn get(
-            &self,
-            _default: OptionalArg<PyObjectRef>,
-            _vm: &VirtualMachine,
-        ) -> PyResult<PyObjectRef> {
-            unimplemented!("ContextVar.get() is currently under construction")
+            zelf: &Py<Self>,
+            default: OptionalArg<PyObjectRef>,
+            vm: &VirtualMachine,
+        ) -> PyResult<Option<PyObjectRef>> {
+            let found = super::CONTEXTS.with(|ctxs| {
+                let ctxs = ctxs.borrow();
+                let ctx = ctxs.last()?;
+                let cached_ptr = zelf.cached.as_ptr();
+                debug_assert!(!cached_ptr.is_null());
+                if let Some(cached) = unsafe { &*cached_ptr } {
+                    if zelf.cached_id.load(Ordering::SeqCst) == ctx.get_id()
+                        && cached.idx + 1 == ctxs.len()
+                    {
+                        return Some(cached.object.clone());
+                    }
+                }
+                let vars = ctx.borrow_vars();
+                let obj = vars.get(zelf)?;
+                zelf.cached_id.store(ctx.get_id(), Ordering::SeqCst);
+
+                // TODO: ensure cached is not changed
+                let _removed = zelf.cached.swap(Some(ContextVarCache {
+                    object: obj.clone(),
+                    idx: ctxs.len() - 1,
+                }));
+
+                Some(obj.clone())
+            });
+
+            let value = if let Some(value) = found {
+                value
+            } else if let Some(default) = default.into_option() {
+                default
+            } else if let Some(default) = zelf.default.as_ref() {
+                default.clone()
+            } else {
+                let msg = zelf.as_object().repr(vm)?;
+                return Err(vm.new_lookup_error(msg.as_str().to_owned()));
+            };
+            Ok(Some(value))
         }
 
         #[pymethod]
-        fn set(&self, _value: PyObjectRef, _vm: &VirtualMachine) -> PyResult<()> {
-            unimplemented!("ContextVar.set() is currently under construction")
+        fn set(
+            zelf: &Py<Self>,
+            value: PyObjectRef,
+            vm: &VirtualMachine,
+        ) -> PyResult<PyRef<ContextToken>> {
+            let ctx = PyContext::current(vm);
+
+            let old_value = ctx.borrow_vars().get(zelf).map(|v| v.to_owned());
+            let token = ContextToken {
+                ctx: ctx.to_owned(),
+                var: zelf.to_owned(),
+                old_value,
+                used: false.into(),
+            };
+
+            // ctx.vars borrow must be released
+            Self::set_inner(zelf, value, vm)?;
+
+            Ok(token.into_ref(&vm.ctx))
         }
 
         #[pymethod]
-        fn reset(
-            _zelf: PyRef<Self>,
-            _token: PyRef<ContextToken>,
-            _vm: &VirtualMachine,
-        ) -> PyResult<()> {
-            unimplemented!("ContextVar.reset() is currently under construction")
+        fn reset(zelf: &Py<Self>, token: PyRef<ContextToken>, vm: &VirtualMachine) -> PyResult<()> {
+            if token.used.get() {
+                let msg = format!("{} has already been used once", token.as_object().repr(vm)?);
+                return Err(vm.new_runtime_error(msg));
+            }
+
+            if !zelf.is(&token.var) {
+                let msg = format!(
+                    "{} was created by a different ContextVar",
+                    token.var.as_object().repr(vm)?
+                );
+                return Err(vm.new_value_error(msg));
+            }
+
+            let ctx = PyContext::current(vm);
+            if !ctx.is(&token.ctx) {
+                let msg = format!(
+                    "{} was created in a different Context",
+                    token.var.as_object().repr(vm)?
+                );
+                return Err(vm.new_value_error(msg));
+            }
+
+            token.used.set(true);
+
+            if let Some(old_value) = &token.old_value {
+                Self::set_inner(zelf, old_value.clone(), vm)?;
+            } else {
+                Self::delete(zelf, vm)?;
+            }
+            Ok(())
         }
 
         #[pyclassmethod(magic)]
@@ -136,73 +479,129 @@ mod _contextvars {
         }
     }
 
-    impl Initializer for ContextVar {
+    #[derive(FromArgs)]
+    struct ContextVarOptions {
+        #[pyarg(positional)]
+        #[allow(dead_code)] // TODO: RUSTPYTHON
+        name: PyStrRef,
+        #[pyarg(any, optional)]
+        #[allow(dead_code)] // TODO: RUSTPYTHON
+        default: OptionalArg<PyObjectRef>,
+    }
+
+    impl Constructor for ContextVar {
         type Args = ContextVarOptions;
+        fn py_new(_cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
+            let var = ContextVar {
+                name: args.name.to_string(),
+                default: args.default.into_option(),
+                cached_id: 0.into(),
+                cached: AtomicCell::new(None),
+                hash: UnsafeCell::new(0),
+            };
+            let py_var = var.into_ref(&vm.ctx);
 
-        fn init(_obj: PyRef<Self>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<()> {
-            unimplemented!("ContextVar.__init__() is currently under construction")
+            unsafe {
+                // SAFETY: py_var is not exposed to python memory model yet
+                *py_var.hash.get() = Self::generate_hash(&py_var, vm)
+            };
+            Ok(py_var.into())
+        }
+    }
+
+    impl std::hash::Hash for ContextVar {
+        #[inline]
+        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+            unsafe { *self.hash.get() }.hash(state)
+        }
+    }
+
+    impl Hashable for ContextVar {
+        #[inline]
+        fn hash(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyHash> {
+            Ok(unsafe { *zelf.hash.get() })
         }
     }
 
     impl Representable for ContextVar {
         #[inline]
-        fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
-            unimplemented!("<ContextVar name={{}} default={{}} at {{}}")
-            // format!(
-            //     "<ContextVar name={} default={:?} at {:#x}>",
-            //     zelf.name.as_str(),
-            //     zelf.default.map_or("", |x| PyStr::from(*x).as_str()),
-            //     zelf.get_id()
-            // )
+        fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
+            // unimplemented!("<ContextVar name={{}} default={{}} at {{}}")
+            Ok(format!(
+                "<ContextVar name={} default={:?} at {:#x}>",
+                zelf.name.as_str(),
+                zelf.default
+                    .as_ref()
+                    .and_then(|default| default.str(vm).ok()),
+                zelf.get_id()
+            ))
         }
     }
 
     #[pyattr]
     #[pyclass(name = "Token")]
     #[derive(Debug, PyPayload)]
-    struct ContextToken {}
-
-    #[derive(FromArgs)]
-    struct ContextTokenOptions {
-        #[pyarg(positional)]
-        #[allow(dead_code)] // TODO: RUSTPYTHON
-        context: PyObjectRef,
-        #[pyarg(positional)]
-        #[allow(dead_code)] // TODO: RUSTPYTHON
-        var: PyObjectRef,
-        #[pyarg(positional)]
-        #[allow(dead_code)] // TODO: RUSTPYTHON
-        old_value: PyObjectRef,
+    struct ContextToken {
+        ctx: PyRef<PyContext>,          // tok_ctx in CPython
+        var: PyRef<ContextVar>,         // tok_var in CPython
+        old_value: Option<PyObjectRef>, // tok_oldval in CPython
+        used: Cell<bool>,
     }
 
-    #[pyclass(with(Initializer, Representable))]
+    unsafe impl Sync for ContextToken {}
+
+    #[pyclass(with(Constructor, Representable))]
     impl ContextToken {
         #[pygetset]
-        fn var(&self, _vm: &VirtualMachine) -> PyObjectRef {
-            unimplemented!("Token.var() is currently under construction")
+        fn var(&self, _vm: &VirtualMachine) -> PyRef<ContextVar> {
+            self.var.clone()
         }
 
         #[pygetset]
         fn old_value(&self, _vm: &VirtualMachine) -> PyObjectRef {
-            unimplemented!("Token.old_value() is currently under construction")
+            match &self.old_value {
+                Some(value) => value.clone(),
+                None => ContextTokenMissing::static_type().to_owned().into(),
+            }
         }
     }
 
-    impl Initializer for ContextToken {
-        type Args = ContextTokenOptions;
+    impl Constructor for ContextToken {
+        type Args = FuncArgs;
 
-        fn init(_obj: PyRef<Self>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<()> {
-            unimplemented!("Token.__init__() is currently under construction")
+        fn slot_new(_cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
+            Err(vm.new_runtime_error("Tokens can only be created by ContextVars".to_owned()))
+        }
+        fn py_new(_cls: PyTypeRef, _args: Self::Args, _vm: &VirtualMachine) -> PyResult {
+            unreachable!()
         }
     }
 
     impl Representable for ContextToken {
         #[inline]
+        fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
+            let used = if zelf.used.get() { " used" } else { "" };
+            let var = Representable::repr_str(&zelf.var, vm)?;
+            let ptr = zelf.as_object().get_id() as *const u8;
+            Ok(format!("<Token{used} var={var} at {ptr:p}>"))
+        }
+    }
+
+    #[pyclass(no_attr, name = "Token.MISSING")]
+    #[derive(Debug, PyPayload)]
+    pub(super) struct ContextTokenMissing {}
+
+    #[pyclass(with(Representable))]
+    impl ContextTokenMissing {}
+
+    impl Representable for ContextTokenMissing {
         fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
-            unimplemented!("<Token {{}}var={{}} at {{}}>")
+            Ok("<Token.MISSING>".to_owned())
         }
     }
 
     #[pyfunction]
-    fn copy_context() {}
+    fn copy_context(vm: &VirtualMachine) -> PyContext {
+        PyContext::current(vm).copy()
+    }
 }
diff --git a/vm/src/object/core.rs b/vm/src/object/core.rs
index bda8aaba2f..68447371fe 100644
--- a/vm/src/object/core.rs
+++ b/vm/src/object/core.rs
@@ -934,6 +934,28 @@ impl<T: PyObjectPayload> Borrow<PyObject> for Py<T> {
     }
 }
 
+impl<T> std::hash::Hash for Py<T>
+where
+    T: std::hash::Hash + PyObjectPayload,
+{
+    #[inline]
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.deref().hash(state)
+    }
+}
+
+impl<T> PartialEq for Py<T>
+where
+    T: PartialEq + PyObjectPayload,
+{
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        self.deref().eq(other.deref())
+    }
+}
+
+impl<T> Eq for Py<T> where T: Eq + PyObjectPayload {}
+
 impl<T> AsRef<PyObject> for Py<T>
 where
     T: PyObjectPayload,
@@ -1089,6 +1111,28 @@ where
     }
 }
 
+impl<T> std::hash::Hash for PyRef<T>
+where
+    T: std::hash::Hash + PyObjectPayload,
+{
+    #[inline]
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.deref().hash(state)
+    }
+}
+
+impl<T> PartialEq for PyRef<T>
+where
+    T: PartialEq + PyObjectPayload,
+{
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        self.deref().eq(other.deref())
+    }
+}
+
+impl<T> Eq for PyRef<T> where T: Eq + PyObjectPayload {}
+
 #[repr(transparent)]
 pub struct PyWeakRef<T: PyObjectPayload> {
     weak: PyRef<PyWeak>,


Test patch:
diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py
index 01b3399327..aa83c6bbd7 100644
--- a/Lib/test/test_context.py
+++ b/Lib/test/test_context.py
@@ -10,7 +10,7 @@
 from test.support import threading_helper
 
 try:
-    from _testcapi import hamt
+    from _testinternalcapi import hamt
 except ImportError:
     hamt = None
 
@@ -42,8 +42,6 @@ def test_context_var_new_1(self):
 
         self.assertNotEqual(hash(c), hash('aaa'))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @isolated_context
     def test_context_var_repr_1(self):
         c = contextvars.ContextVar('a')
@@ -101,8 +99,6 @@ def test_context_typerrors_1(self):
         with self.assertRaisesRegex(TypeError, 'ContextVar key was expected'):
             ctx.get(1)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_context_get_context_1(self):
         ctx = contextvars.copy_context()
         self.assertIsInstance(ctx, contextvars.Context)
@@ -115,8 +111,6 @@ def test_context_run_1(self):
         with self.assertRaisesRegex(TypeError, 'missing 1 required'):
             ctx.run()
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_context_run_2(self):
         ctx = contextvars.Context()
 
@@ -145,8 +139,6 @@ def func(*args, **kwargs):
                 ((11, 'bar'), {'spam': 'foo'}))
             self.assertEqual(a, {})
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_context_run_3(self):
         ctx = contextvars.Context()
 
@@ -187,8 +179,6 @@ def func1():
         self.assertEqual(returned_ctx[var], 'spam')
         self.assertIn(var, returned_ctx)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_context_run_5(self):
         ctx = contextvars.Context()
         var = contextvars.ContextVar('var')
@@ -203,8 +193,6 @@ def func():
 
         self.assertIsNone(var.get(None))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_context_run_6(self):
         ctx = contextvars.Context()
         c = contextvars.ContextVar('a', default=0)
@@ -219,8 +207,6 @@ def fun():
 
         ctx.run(fun)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_context_run_7(self):
         ctx = contextvars.Context()
 
@@ -286,8 +272,6 @@ def test_context_getset_1(self):
         self.assertEqual(len(ctx2), 0)
         self.assertEqual(list(ctx2), [])
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @isolated_context
     def test_context_getset_2(self):
         v1 = contextvars.ContextVar('v1')
@@ -297,8 +281,6 @@ def test_context_getset_2(self):
         with self.assertRaisesRegex(ValueError, 'by a different'):
             v2.reset(t1)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @isolated_context
     def test_context_getset_3(self):
         c = contextvars.ContextVar('c', default=42)
@@ -324,8 +306,6 @@ def fun():
 
         ctx.run(fun)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @isolated_context
     def test_context_getset_4(self):
         c = contextvars.ContextVar('c', default=42)
@@ -378,8 +358,6 @@ def ctx2_fun():
 
         ctx1.run(ctx1_fun)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @isolated_context
     @threading_helper.requires_working_threading()
     def test_context_threads_1(self):
@@ -470,7 +448,7 @@ class EqError(Exception):
     pass
 
 
-@unittest.skipIf(hamt is None, '_testcapi lacks "hamt()" function')
+@unittest.skipIf(hamt is None, '_testinternalcapi.hamt() not available')
 class HamtTest(unittest.TestCase):
 
     def test_hashkey_helper_1(self):