1use super::{PyStr, PyType, PyTypeRef};
8use crate::{
9 AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
10 builtins::function::PyCell,
11 class::PyClassImpl,
12 common::lock::PyRwLock,
13 function::{FuncArgs, IntoFuncArgs, OptionalArg},
14 types::{Callable, Constructor, GetAttr, GetDescriptor, Initializer, Representable},
15};
16
17#[pyclass(module = false, name = "super", traverse)]
18#[derive(Debug)]
19pub struct PySuper {
20 inner: PyRwLock<PySuperInner>,
21}
22
23#[derive(Debug, Traverse)]
24struct PySuperInner {
25 typ: PyTypeRef,
26 obj: Option<(PyObjectRef, PyTypeRef)>,
27}
28
29impl PySuperInner {
30 fn new(typ: PyTypeRef, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
31 let obj = if vm.is_none(&obj) {
32 None
33 } else {
34 let obj_type = super_check(typ.clone(), obj.clone(), vm)?;
35 Some((obj, obj_type))
36 };
37 Ok(Self { typ, obj })
38 }
39}
40
41impl PyPayload for PySuper {
42 #[inline]
43 fn class(ctx: &Context) -> &'static Py<PyType> {
44 ctx.types.super_type
45 }
46}
47
48impl Constructor for PySuper {
49 type Args = FuncArgs;
50
51 fn py_new(_cls: &Py<PyType>, _args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
52 Ok(Self {
53 inner: PyRwLock::new(PySuperInner::new(
54 vm.ctx.types.object_type.to_owned(), vm.ctx.none(),
56 vm,
57 )?),
58 })
59 }
60}
61
62#[derive(FromArgs)]
63pub struct InitArgs {
64 #[pyarg(positional, optional, error_msg = "super() argument 1 must be a type")]
65 py_type: OptionalArg<PyTypeRef>,
66 #[pyarg(positional, optional)]
67 py_obj: OptionalArg<PyObjectRef>,
68}
69
70impl Initializer for PySuper {
71 type Args = InitArgs;
72
73 fn init(
74 zelf: PyRef<Self>,
75 Self::Args { py_type, py_obj }: Self::Args,
76 vm: &VirtualMachine,
77 ) -> PyResult<()> {
78 let (typ, obj) = if let OptionalArg::Present(ty) = py_type {
80 (ty, py_obj.unwrap_or_none(vm))
81 } else {
82 let frame = vm
83 .current_frame()
84 .ok_or_else(|| vm.new_runtime_error("super(): no current frame"))?;
85
86 if frame.code.arg_count == 0 {
87 return Err(vm.new_runtime_error("super(): no arguments"));
88 }
89 use rustpython_compiler_core::bytecode::CO_FAST_CELL;
91 let obj = unsafe { frame.fastlocals() }[0]
92 .clone()
93 .and_then(|val| {
94 if frame
96 .code
97 .localspluskinds
98 .first()
99 .is_some_and(|&k| k & CO_FAST_CELL != 0)
100 {
101 val.downcast_ref::<PyCell>().and_then(|c| c.get())
102 } else {
103 Some(val)
104 }
105 })
106 .ok_or_else(|| vm.new_runtime_error("super(): arg[0] deleted"))?;
107
108 let mut typ = None;
109 let nlocalsplus = frame.code.localspluskinds.len();
111 let nfrees = frame.code.freevars.len();
112 let free_start = nlocalsplus - nfrees;
113 for (i, var) in frame.code.freevars.iter().enumerate() {
114 if var.as_bytes() == b"__class__" {
115 let class = frame
116 .get_cell_contents(free_start + i)
117 .ok_or_else(|| vm.new_runtime_error("super(): empty __class__ cell"))?;
118 typ = Some(class.downcast().map_err(|o| {
119 vm.new_type_error(format!(
120 "super(): __class__ is not a type ({})",
121 o.class().name()
122 ))
123 })?);
124 break;
125 }
126 }
127 let typ = typ.ok_or_else(|| {
128 vm.new_type_error(
129 "super must be called with 1 argument or from inside class method",
130 )
131 })?;
132
133 (typ, obj)
134 };
135
136 let inner = PySuperInner::new(typ, obj, vm)?;
137 *zelf.inner.write() = inner;
138
139 Ok(())
140 }
141}
142
143#[pyclass(with(GetAttr, GetDescriptor, Constructor, Initializer, Representable))]
144impl PySuper {
145 #[pygetset]
146 fn __thisclass__(&self) -> PyTypeRef {
147 self.inner.read().typ.clone()
148 }
149
150 #[pygetset]
151 fn __self_class__(&self) -> Option<PyTypeRef> {
152 Some(self.inner.read().obj.as_ref()?.1.clone())
153 }
154
155 #[pygetset]
156 fn __self__(&self) -> Option<PyObjectRef> {
157 Some(self.inner.read().obj.as_ref()?.0.clone())
158 }
159}
160
161impl GetAttr for PySuper {
162 fn getattro(zelf: &Py<Self>, name: &Py<PyStr>, vm: &VirtualMachine) -> PyResult {
163 let skip = |zelf: &Py<Self>, name| zelf.as_object().generic_getattr(name, vm);
164 let (obj, start_type): (PyObjectRef, PyTypeRef) = match &zelf.inner.read().obj {
165 Some(o) => o.clone(),
166 None => return skip(zelf, name),
167 };
168 if name.as_bytes() == b"__class__" {
172 return skip(zelf, name);
173 }
174
175 if let Some(name) = vm.ctx.interned_str(name) {
176 let mro: Vec<PyRef<PyType>> = start_type.mro_map_collect(|x| x.to_owned());
178 let mro: Vec<_> = mro
179 .iter()
180 .skip_while(|cls| !cls.is(&zelf.inner.read().typ))
181 .skip(1) .collect();
183 for cls in &mro {
184 if let Some(descr) = cls.get_direct_attr(name) {
185 return vm
186 .call_get_descriptor_specific(
187 &descr,
188 if obj.is(&start_type) { None } else { Some(obj) },
190 Some(start_type.as_object().to_owned()),
191 )
192 .unwrap_or(Ok(descr));
193 }
194 }
195 }
196 skip(zelf, name)
197 }
198}
199
200impl GetDescriptor for PySuper {
201 fn descr_get(
202 zelf_obj: PyObjectRef,
203 obj: Option<PyObjectRef>,
204 _cls: Option<PyObjectRef>,
205 vm: &VirtualMachine,
206 ) -> PyResult {
207 let (zelf, obj) = Self::_unwrap(&zelf_obj, obj, vm)?;
208 if vm.is_none(&obj) || zelf.inner.read().obj.is_some() {
209 return Ok(zelf_obj);
210 }
211 let zelf_class = zelf.as_object().class();
212 if zelf_class.is(vm.ctx.types.super_type) {
213 let typ = zelf.inner.read().typ.clone();
214 Ok(Self {
215 inner: PyRwLock::new(PySuperInner::new(typ, obj, vm)?),
216 }
217 .into_ref(&vm.ctx)
218 .into())
219 } else {
220 let (obj, typ) = {
221 let lock = zelf.inner.read();
222 let obj = lock.obj.as_ref().map(|(o, _)| o.to_owned());
223 let typ = lock.typ.clone();
224 (obj, typ)
225 };
226 let obj = vm.unwrap_or_none(obj);
227 PyType::call(zelf.class(), (typ, obj).into_args(vm), vm)
228 }
229 }
230}
231
232impl Representable for PySuper {
233 #[inline]
234 fn repr_str(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
235 let type_name = zelf.inner.read().typ.name().to_owned();
236 let obj = zelf.inner.read().obj.clone();
237 let repr = match obj {
238 Some((_, ref ty)) => {
239 format!("<super: <class '{}'>, <{} object>>", &type_name, ty.name())
240 }
241 None => format!("<super: <class '{type_name}'>, NULL>"),
242 };
243 Ok(repr)
244 }
245}
246
247fn super_check(ty: PyTypeRef, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyTypeRef> {
248 let typ = match obj.clone().downcast::<PyType>() {
249 Ok(cls) if cls.fast_issubclass(&ty) => return Ok(cls),
250 Ok(cls) => Some(cls),
251 Err(_) => None,
252 };
253
254 if obj.fast_isinstance(&ty) {
255 return Ok(obj.class().to_owned());
256 }
257
258 let class_attr = obj.get_attr("__class__", vm)?;
259 if let Ok(cls) = class_attr.downcast::<PyType>()
260 && !cls.is(&ty)
261 && cls.fast_issubclass(&ty)
262 {
263 return Ok(cls);
264 }
265
266 let (type_or_instance, obj_str) = match typ {
267 Some(t) => ("type", t.name().to_owned()),
268 None => ("instance of", obj.class().name().to_owned()),
269 };
270
271 Err(vm.new_type_error(format!(
272 "super(type, obj): obj ({} {}) is not an instance or subtype of type ({}).",
273 type_or_instance,
274 obj_str,
275 ty.name(),
276 )))
277}
278
279pub fn init(context: &'static Context) {
280 let super_type = &context.types.super_type;
281 PySuper::extend_class(context, super_type);
282
283 let super_doc = "super() -> same as super(__class__, <first argument>)\n\
284 super(type) -> unbound super object\n\
285 super(type, obj) -> bound super object; requires isinstance(obj, type)\n\
286 super(type, type2) -> bound super object; requires issubclass(type2, type)\n\
287 Typical use to call a cooperative superclass method:\n\
288 class C(B):\n \
289 def meth(self, arg):\n \
290 super().meth(arg)\n\
291 This works for class methods too:\n\
292 class C(B):\n \
293 @classmethod\n \
294 def cmeth(cls, arg):\n \
295 super().cmeth(arg)\n";
296
297 extend_class!(context, super_type, {
298 "__doc__" => context.new_str(super_doc),
299 });
300}