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