rustpython_vm/builtins/
genericalias.rs

1use once_cell::sync::Lazy;
2
3use super::type_;
4use crate::{
5    atomic_func,
6    builtins::{PyList, PyStr, PyTuple, PyTupleRef, PyType, PyTypeRef},
7    class::PyClassImpl,
8    common::hash,
9    convert::ToPyObject,
10    function::{FuncArgs, PyComparisonValue},
11    protocol::{PyMappingMethods, PyNumberMethods},
12    types::{
13        AsMapping, AsNumber, Callable, Comparable, Constructor, GetAttr, Hashable, PyComparisonOp,
14        Representable,
15    },
16    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
17    VirtualMachine,
18};
19use std::fmt;
20
21static ATTR_EXCEPTIONS: [&str; 8] = [
22    "__origin__",
23    "__args__",
24    "__parameters__",
25    "__mro_entries__",
26    "__reduce_ex__", // needed so we don't look up object.__reduce_ex__
27    "__reduce__",
28    "__copy__",
29    "__deepcopy__",
30];
31
32#[pyclass(module = "types", name = "GenericAlias")]
33pub struct PyGenericAlias {
34    origin: PyTypeRef,
35    args: PyTupleRef,
36    parameters: PyTupleRef,
37}
38
39impl fmt::Debug for PyGenericAlias {
40    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
41        f.write_str("GenericAlias")
42    }
43}
44
45impl PyPayload for PyGenericAlias {
46    fn class(ctx: &Context) -> &'static Py<PyType> {
47        ctx.types.generic_alias_type
48    }
49}
50
51impl Constructor for PyGenericAlias {
52    type Args = FuncArgs;
53
54    fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
55        if !args.kwargs.is_empty() {
56            return Err(vm.new_type_error("GenericAlias() takes no keyword arguments".to_owned()));
57        }
58        let (origin, arguments): (_, PyObjectRef) = args.bind(vm)?;
59        PyGenericAlias::new(origin, arguments, vm)
60            .into_ref_with_type(vm, cls)
61            .map(Into::into)
62    }
63}
64
65#[pyclass(
66    with(
67        AsNumber,
68        AsMapping,
69        Callable,
70        Comparable,
71        Constructor,
72        GetAttr,
73        Hashable,
74        Representable
75    ),
76    flags(BASETYPE)
77)]
78impl PyGenericAlias {
79    pub fn new(origin: PyTypeRef, args: PyObjectRef, vm: &VirtualMachine) -> Self {
80        let args = if let Ok(tuple) = args.try_to_ref::<PyTuple>(vm) {
81            tuple.to_owned()
82        } else {
83            PyTuple::new_ref(vec![args], &vm.ctx)
84        };
85
86        let parameters = make_parameters(&args, vm);
87        Self {
88            origin,
89            args,
90            parameters,
91        }
92    }
93
94    fn repr(&self, vm: &VirtualMachine) -> PyResult<String> {
95        fn repr_item(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<String> {
96            if obj.is(&vm.ctx.ellipsis) {
97                return Ok("...".to_string());
98            }
99
100            if vm
101                .get_attribute_opt(obj.clone(), identifier!(vm, __origin__))?
102                .is_some()
103                && vm
104                    .get_attribute_opt(obj.clone(), identifier!(vm, __args__))?
105                    .is_some()
106            {
107                return Ok(obj.repr(vm)?.as_str().to_string());
108            }
109
110            match (
111                vm.get_attribute_opt(obj.clone(), identifier!(vm, __qualname__))?
112                    .and_then(|o| o.downcast_ref::<PyStr>().map(|n| n.as_str().to_string())),
113                vm.get_attribute_opt(obj.clone(), identifier!(vm, __module__))?
114                    .and_then(|o| o.downcast_ref::<PyStr>().map(|m| m.as_str().to_string())),
115            ) {
116                (None, _) | (_, None) => Ok(obj.repr(vm)?.as_str().to_string()),
117                (Some(qualname), Some(module)) => Ok(if module == "builtins" {
118                    qualname
119                } else {
120                    format!("{module}.{qualname}")
121                }),
122            }
123        }
124
125        Ok(format!(
126            "{}[{}]",
127            repr_item(self.origin.clone().into(), vm)?,
128            if self.args.len() == 0 {
129                "()".to_owned()
130            } else {
131                self.args
132                    .iter()
133                    .map(|o| repr_item(o.clone(), vm))
134                    .collect::<PyResult<Vec<_>>>()?
135                    .join(", ")
136            }
137        ))
138    }
139
140    #[pygetset(magic)]
141    fn parameters(&self) -> PyObjectRef {
142        self.parameters.clone().into()
143    }
144
145    #[pygetset(magic)]
146    fn args(&self) -> PyObjectRef {
147        self.args.clone().into()
148    }
149
150    #[pygetset(magic)]
151    fn origin(&self) -> PyObjectRef {
152        self.origin.clone().into()
153    }
154
155    #[pymethod(magic)]
156    fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
157        let new_args = subs_parameters(
158            |vm| self.repr(vm),
159            self.args.clone(),
160            self.parameters.clone(),
161            needle,
162            vm,
163        )?;
164
165        Ok(
166            PyGenericAlias::new(self.origin.clone(), new_args.to_pyobject(vm), vm)
167                .into_pyobject(vm),
168        )
169    }
170
171    #[pymethod(magic)]
172    fn dir(&self, vm: &VirtualMachine) -> PyResult<PyList> {
173        let dir = vm.dir(Some(self.origin()))?;
174        for exc in &ATTR_EXCEPTIONS {
175            if !dir.contains((*exc).to_pyobject(vm), vm)? {
176                dir.append((*exc).to_pyobject(vm));
177            }
178        }
179        Ok(dir)
180    }
181
182    #[pymethod(magic)]
183    fn reduce(zelf: &Py<Self>, vm: &VirtualMachine) -> (PyTypeRef, (PyTypeRef, PyTupleRef)) {
184        (
185            vm.ctx.types.generic_alias_type.to_owned(),
186            (zelf.origin.clone(), zelf.args.clone()),
187        )
188    }
189
190    #[pymethod(magic)]
191    fn mro_entries(&self, _bases: PyObjectRef, vm: &VirtualMachine) -> PyTupleRef {
192        PyTuple::new_ref(vec![self.origin()], &vm.ctx)
193    }
194
195    #[pymethod(magic)]
196    fn instancecheck(_zelf: PyRef<Self>, _obj: PyObjectRef, vm: &VirtualMachine) -> PyResult {
197        Err(vm
198            .new_type_error("isinstance() argument 2 cannot be a parameterized generic".to_owned()))
199    }
200
201    #[pymethod(magic)]
202    fn subclasscheck(_zelf: PyRef<Self>, _obj: PyObjectRef, vm: &VirtualMachine) -> PyResult {
203        Err(vm
204            .new_type_error("issubclass() argument 2 cannot be a parameterized generic".to_owned()))
205    }
206
207    #[pymethod(magic)]
208    fn ror(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
209        type_::or_(other, zelf, vm)
210    }
211
212    #[pymethod(magic)]
213    fn or(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
214        type_::or_(zelf, other, vm)
215    }
216}
217
218pub(crate) fn is_typevar(obj: &PyObjectRef, vm: &VirtualMachine) -> bool {
219    let class = obj.class();
220    "TypeVar" == &*class.slot_name()
221        && class
222            .get_attr(identifier!(vm, __module__))
223            .and_then(|o| o.downcast_ref::<PyStr>().map(|s| s.as_str() == "typing"))
224            .unwrap_or(false)
225}
226
227pub(crate) fn make_parameters(args: &Py<PyTuple>, vm: &VirtualMachine) -> PyTupleRef {
228    let mut parameters: Vec<PyObjectRef> = Vec::with_capacity(args.len());
229    for arg in args {
230        if is_typevar(arg, vm) {
231            if !parameters.iter().any(|param| param.is(arg)) {
232                parameters.push(arg.clone());
233            }
234        } else if let Ok(obj) = arg.get_attr(identifier!(vm, __parameters__), vm) {
235            if let Ok(sub_params) = obj.try_to_ref::<PyTuple>(vm) {
236                for sub_param in sub_params {
237                    if !parameters.iter().any(|param| param.is(sub_param)) {
238                        parameters.push(sub_param.clone());
239                    }
240                }
241            }
242        }
243    }
244    parameters.shrink_to_fit();
245
246    PyTuple::new_ref(parameters, &vm.ctx)
247}
248
249#[inline]
250fn tuple_index(tuple: &PyTupleRef, item: &PyObjectRef) -> Option<usize> {
251    tuple.iter().position(|element| element.is(item))
252}
253
254fn subs_tvars(
255    obj: PyObjectRef,
256    params: &PyTupleRef,
257    argitems: &[PyObjectRef],
258    vm: &VirtualMachine,
259) -> PyResult {
260    obj.get_attr(identifier!(vm, __parameters__), vm)
261        .ok()
262        .and_then(|sub_params| {
263            PyTupleRef::try_from_object(vm, sub_params)
264                .ok()
265                .and_then(|sub_params| {
266                    if sub_params.len() > 0 {
267                        let sub_args = sub_params
268                            .iter()
269                            .map(|arg| {
270                                if let Some(idx) = tuple_index(params, arg) {
271                                    argitems[idx].clone()
272                                } else {
273                                    arg.clone()
274                                }
275                            })
276                            .collect::<Vec<_>>();
277                        let sub_args: PyObjectRef = PyTuple::new_ref(sub_args, &vm.ctx).into();
278                        Some(obj.get_item(&*sub_args, vm))
279                    } else {
280                        None
281                    }
282                })
283        })
284        .unwrap_or(Ok(obj))
285}
286
287pub fn subs_parameters<F: Fn(&VirtualMachine) -> PyResult<String>>(
288    repr: F,
289    args: PyTupleRef,
290    parameters: PyTupleRef,
291    needle: PyObjectRef,
292    vm: &VirtualMachine,
293) -> PyResult<PyTupleRef> {
294    let num_params = parameters.len();
295    if num_params == 0 {
296        return Err(vm.new_type_error(format!("There are no type variables left in {}", repr(vm)?)));
297    }
298
299    let items = needle.try_to_ref::<PyTuple>(vm);
300    let arg_items = match items {
301        Ok(tuple) => tuple.as_slice(),
302        Err(_) => std::slice::from_ref(&needle),
303    };
304
305    let num_items = arg_items.len();
306    if num_params != num_items {
307        let plural = if num_items > num_params {
308            "many"
309        } else {
310            "few"
311        };
312        return Err(vm.new_type_error(format!("Too {} arguments for {}", plural, repr(vm)?)));
313    }
314
315    let new_args = args
316        .iter()
317        .map(|arg| {
318            if is_typevar(arg, vm) {
319                let idx = tuple_index(&parameters, arg).unwrap();
320                Ok(arg_items[idx].clone())
321            } else {
322                subs_tvars(arg.clone(), &parameters, arg_items, vm)
323            }
324        })
325        .collect::<PyResult<Vec<_>>>()?;
326
327    Ok(PyTuple::new_ref(new_args, &vm.ctx))
328}
329
330impl AsMapping for PyGenericAlias {
331    fn as_mapping() -> &'static PyMappingMethods {
332        static AS_MAPPING: Lazy<PyMappingMethods> = Lazy::new(|| PyMappingMethods {
333            subscript: atomic_func!(|mapping, needle, vm| {
334                PyGenericAlias::mapping_downcast(mapping).getitem(needle.to_owned(), vm)
335            }),
336            ..PyMappingMethods::NOT_IMPLEMENTED
337        });
338        &AS_MAPPING
339    }
340}
341
342impl AsNumber for PyGenericAlias {
343    fn as_number() -> &'static PyNumberMethods {
344        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
345            or: Some(|a, b, vm| Ok(PyGenericAlias::or(a.to_owned(), b.to_owned(), vm))),
346            ..PyNumberMethods::NOT_IMPLEMENTED
347        };
348        &AS_NUMBER
349    }
350}
351
352impl Callable for PyGenericAlias {
353    type Args = FuncArgs;
354    fn call(zelf: &Py<Self>, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
355        PyType::call(&zelf.origin, args, vm).map(|obj| {
356            if let Err(exc) = obj.set_attr(identifier!(vm, __orig_class__), zelf.to_owned(), vm) {
357                if !exc.fast_isinstance(vm.ctx.exceptions.attribute_error)
358                    && !exc.fast_isinstance(vm.ctx.exceptions.type_error)
359                {
360                    return Err(exc);
361                }
362            }
363            Ok(obj)
364        })?
365    }
366}
367
368impl Comparable for PyGenericAlias {
369    fn cmp(
370        zelf: &Py<Self>,
371        other: &PyObject,
372        op: PyComparisonOp,
373        vm: &VirtualMachine,
374    ) -> PyResult<PyComparisonValue> {
375        op.eq_only(|| {
376            let other = class_or_notimplemented!(Self, other);
377            Ok(PyComparisonValue::Implemented(
378                if !zelf
379                    .origin()
380                    .rich_compare_bool(&other.origin(), PyComparisonOp::Eq, vm)?
381                {
382                    false
383                } else {
384                    zelf.args()
385                        .rich_compare_bool(&other.args(), PyComparisonOp::Eq, vm)?
386                },
387            ))
388        })
389    }
390}
391
392impl Hashable for PyGenericAlias {
393    #[inline]
394    fn hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<hash::PyHash> {
395        Ok(zelf.origin.as_object().hash(vm)? ^ zelf.args.as_object().hash(vm)?)
396    }
397}
398
399impl GetAttr for PyGenericAlias {
400    fn getattro(zelf: &Py<Self>, attr: &Py<PyStr>, vm: &VirtualMachine) -> PyResult {
401        for exc in ATTR_EXCEPTIONS.iter() {
402            if *(*exc) == attr.to_string() {
403                return zelf.as_object().generic_getattr(attr, vm);
404            }
405        }
406        zelf.origin().get_attr(attr, vm)
407    }
408}
409
410impl Representable for PyGenericAlias {
411    #[inline]
412    fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
413        zelf.repr(vm)
414    }
415}
416
417pub fn init(context: &Context) {
418    let generic_alias_type = &context.types.generic_alias_type;
419    PyGenericAlias::extend_class(context, generic_alias_type);
420}
rustpython_vm/builtins/genericalias.rs

rustpython_vm/builtins/
genericalias.rs
