rustpython_vm/stdlib/
typevar.rs

1// spell-checker:ignore typevarobject funcobj
2
3pub use typevar::*;
4
5#[pymodule(sub)]
6pub(crate) mod typevar {
7    use crate::{
8        AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
9        builtins::{PyTuple, PyTupleRef, PyType, PyTypeRef, make_union},
10        common::lock::PyMutex,
11        function::{FuncArgs, PyComparisonValue},
12        protocol::PyNumberMethods,
13        stdlib::_typing::{call_typing_func_object, decl::const_evaluator_alloc},
14        types::{AsNumber, Comparable, Constructor, Iterable, PyComparisonOp, Representable},
15    };
16
17    fn type_check(arg: PyObjectRef, msg: &str, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
18        // Calling typing.py here leads to bootstrapping problems
19        if vm.is_none(&arg) {
20            return Ok(arg.class().to_owned().into());
21        }
22        let message_str: PyObjectRef = vm.ctx.new_str(msg).into();
23        call_typing_func_object(vm, "_type_check", (arg, message_str))
24    }
25
26    fn variance_repr(
27        name: &str,
28        infer_variance: bool,
29        covariant: bool,
30        contravariant: bool,
31    ) -> String {
32        if infer_variance {
33            return name.to_owned();
34        }
35        let prefix = if covariant {
36            '+'
37        } else if contravariant {
38            '-'
39        } else {
40            '~'
41        };
42        format!("{prefix}{name}")
43    }
44
45    /// Get the module of the caller frame, similar to CPython's caller() function.
46    /// Returns the module name or None if not found.
47    ///
48    /// Note: CPython's implementation (in typevarobject.c) gets the module from the
49    /// frame's function object using PyFunction_GetModule(f->f_funcobj). However,
50    /// RustPython's Frame doesn't store a reference to the function object, so we
51    /// get the module name from the frame's globals dictionary instead.
52    fn caller(vm: &VirtualMachine) -> Option<PyObjectRef> {
53        let frame = vm.current_frame()?;
54
55        // In RustPython, we get the module name from frame's globals
56        // This is similar to CPython's sys._getframe().f_globals.get('__name__')
57        frame.globals.get_item("__name__", vm).ok()
58    }
59
60    /// Set __module__ attribute for an object based on the caller's module.
61    /// This follows CPython's behavior for TypeVar and similar objects.
62    fn set_module_from_caller(obj: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
63        // Note: CPython gets module from frame->f_funcobj, but RustPython's Frame
64        // architecture is different - we use globals['__name__'] instead
65        let module_value: PyObjectRef = if let Some(module_name) = caller(vm) {
66            // Special handling for certain module names
67            if let Ok(name_str) = module_name.str(vm)
68                && let Some(name) = name_str.to_str()
69                && (name == "builtins" || name.starts_with('<'))
70            {
71                return Ok(());
72            }
73            module_name
74        } else {
75            vm.ctx.none()
76        };
77        obj.set_attr("__module__", module_value, vm)?;
78        Ok(())
79    }
80
81    #[pyattr]
82    #[pyclass(name = "TypeVar", module = "typing")]
83    #[derive(Debug, PyPayload)]
84    #[allow(dead_code)]
85    pub struct TypeVar {
86        name: PyObjectRef, // TODO PyStrRef?
87        bound: PyMutex<PyObjectRef>,
88        evaluate_bound: PyObjectRef,
89        constraints: PyMutex<PyObjectRef>,
90        evaluate_constraints: PyObjectRef,
91        default_value: PyMutex<PyObjectRef>,
92        evaluate_default: PyMutex<PyObjectRef>,
93        covariant: bool,
94        contravariant: bool,
95        infer_variance: bool,
96    }
97    #[pyclass(
98        flags(HAS_DICT, HAS_WEAKREF),
99        with(AsNumber, Constructor, Representable)
100    )]
101    impl TypeVar {
102        #[pymethod]
103        fn __mro_entries__(&self, _bases: PyObjectRef, vm: &VirtualMachine) -> PyResult {
104            Err(vm.new_type_error("Cannot subclass an instance of TypeVar"))
105        }
106
107        #[pygetset]
108        fn __name__(&self) -> PyObjectRef {
109            self.name.clone()
110        }
111
112        #[pygetset]
113        fn __constraints__(&self, vm: &VirtualMachine) -> PyResult {
114            let mut constraints = self.constraints.lock();
115            if !vm.is_none(&constraints) {
116                return Ok(constraints.clone());
117            }
118            let r = if !vm.is_none(&self.evaluate_constraints) {
119                *constraints = self.evaluate_constraints.call((1i32,), vm)?;
120                constraints.clone()
121            } else {
122                vm.ctx.empty_tuple.clone().into()
123            };
124            Ok(r)
125        }
126
127        #[pygetset]
128        fn __bound__(&self, vm: &VirtualMachine) -> PyResult {
129            let mut bound = self.bound.lock();
130            if !vm.is_none(&bound) {
131                return Ok(bound.clone());
132            }
133            let r = if !vm.is_none(&self.evaluate_bound) {
134                *bound = self.evaluate_bound.call((1i32,), vm)?;
135                bound.clone()
136            } else {
137                vm.ctx.none()
138            };
139            Ok(r)
140        }
141
142        #[pygetset]
143        const fn __covariant__(&self) -> bool {
144            self.covariant
145        }
146
147        #[pygetset]
148        const fn __contravariant__(&self) -> bool {
149            self.contravariant
150        }
151
152        #[pygetset]
153        const fn __infer_variance__(&self) -> bool {
154            self.infer_variance
155        }
156
157        #[pygetset]
158        fn __default__(&self, vm: &VirtualMachine) -> PyResult {
159            {
160                let default_value = self.default_value.lock();
161                if !default_value.is(&vm.ctx.typing_no_default) {
162                    return Ok(default_value.clone());
163                }
164            }
165            let evaluator = self.evaluate_default.lock().clone();
166            if !vm.is_none(&evaluator) {
167                let result = evaluator.call((1i32,), vm)?;
168                *self.default_value.lock() = result.clone();
169                Ok(result)
170            } else {
171                Ok(vm.ctx.typing_no_default.clone().into())
172            }
173        }
174
175        #[pygetset]
176        fn evaluate_bound(&self, vm: &VirtualMachine) -> PyResult {
177            if !vm.is_none(&self.evaluate_bound) {
178                return Ok(self.evaluate_bound.clone());
179            }
180            let bound = self.bound.lock();
181            if !vm.is_none(&bound) {
182                return Ok(const_evaluator_alloc(bound.clone(), vm));
183            }
184            Ok(vm.ctx.none())
185        }
186
187        #[pygetset]
188        fn evaluate_constraints(&self, vm: &VirtualMachine) -> PyResult {
189            if !vm.is_none(&self.evaluate_constraints) {
190                return Ok(self.evaluate_constraints.clone());
191            }
192            let constraints = self.constraints.lock();
193            if !vm.is_none(&constraints) {
194                return Ok(const_evaluator_alloc(constraints.clone(), vm));
195            }
196            Ok(vm.ctx.none())
197        }
198
199        #[pygetset]
200        fn evaluate_default(&self, vm: &VirtualMachine) -> PyResult {
201            let evaluator = self.evaluate_default.lock().clone();
202            if !vm.is_none(&evaluator) {
203                return Ok(evaluator);
204            }
205            let default_value = self.default_value.lock().clone();
206            if !default_value.is(&vm.ctx.typing_no_default) {
207                return Ok(const_evaluator_alloc(default_value, vm));
208            }
209            Ok(vm.ctx.none())
210        }
211
212        #[pymethod]
213        fn __typing_subst__(
214            zelf: crate::PyRef<Self>,
215            arg: PyObjectRef,
216            vm: &VirtualMachine,
217        ) -> PyResult {
218            let self_obj: PyObjectRef = zelf.into();
219            call_typing_func_object(vm, "_typevar_subst", (self_obj, arg))
220        }
221
222        #[pymethod]
223        fn __reduce__(&self) -> PyObjectRef {
224            self.name.clone()
225        }
226
227        #[pymethod]
228        fn has_default(&self, vm: &VirtualMachine) -> bool {
229            if !vm.is_none(&self.evaluate_default.lock()) {
230                return true;
231            }
232            let default_value = self.default_value.lock();
233            // Check if default_value is not NoDefault
234            !default_value.is(&vm.ctx.typing_no_default)
235        }
236
237        #[pymethod]
238        fn __typing_prepare_subst__(
239            zelf: crate::PyRef<Self>,
240            alias: PyObjectRef,
241            args: PyObjectRef,
242            vm: &VirtualMachine,
243        ) -> PyResult {
244            // Convert args to tuple if needed
245            let args_tuple =
246                if let Ok(tuple) = args.try_to_ref::<rustpython_vm::builtins::PyTuple>(vm) {
247                    tuple
248                } else {
249                    return Ok(args);
250                };
251
252            // Get alias.__parameters__
253            let parameters = alias.get_attr(identifier!(vm, __parameters__), vm)?;
254            let params_tuple: PyTupleRef = parameters.try_into_value(vm)?;
255
256            // Find our index in parameters
257            let self_obj: PyObjectRef = zelf.to_owned().into();
258            let param_index = params_tuple.iter().position(|p| p.is(&self_obj));
259
260            if let Some(index) = param_index {
261                // Check if we have enough arguments
262                if args_tuple.len() <= index && zelf.has_default(vm) {
263                    // Need to add default value
264                    let mut new_args: Vec<PyObjectRef> = args_tuple.iter().cloned().collect();
265
266                    // Add default value at the correct position
267                    while new_args.len() <= index {
268                        // For the current parameter, add its default
269                        if new_args.len() == index {
270                            let default_val = zelf.__default__(vm)?;
271                            new_args.push(default_val);
272                        } else {
273                            // This shouldn't happen in well-formed code
274                            break;
275                        }
276                    }
277
278                    return Ok(rustpython_vm::builtins::PyTuple::new_ref(new_args, &vm.ctx).into());
279                }
280            }
281
282            // No changes needed
283            Ok(args)
284        }
285    }
286
287    impl Representable for TypeVar {
288        #[inline(always)]
289        fn repr_str(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
290            let name = zelf.name.str_utf8(vm)?;
291            Ok(variance_repr(
292                name.as_str(),
293                zelf.infer_variance,
294                zelf.covariant,
295                zelf.contravariant,
296            ))
297        }
298    }
299
300    impl AsNumber for TypeVar {
301        fn as_number() -> &'static PyNumberMethods {
302            static AS_NUMBER: PyNumberMethods = PyNumberMethods {
303                or: Some(|a, b, vm| {
304                    let args = PyTuple::new_ref(vec![a.to_owned(), b.to_owned()], &vm.ctx);
305                    make_union(&args, vm)
306                }),
307                ..PyNumberMethods::NOT_IMPLEMENTED
308            };
309            &AS_NUMBER
310        }
311    }
312
313    impl Constructor for TypeVar {
314        type Args = FuncArgs;
315
316        fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
317            let typevar = <Self as Constructor>::py_new(&cls, args, vm)?;
318            let obj = typevar.into_ref_with_type(vm, cls)?;
319            let obj_ref: PyObjectRef = obj.into();
320            set_module_from_caller(&obj_ref, vm)?;
321            Ok(obj_ref)
322        }
323
324        fn py_new(_cls: &Py<PyType>, args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
325            let mut kwargs = args.kwargs;
326            // Parse arguments manually
327            let (name, constraints) = if args.args.is_empty() {
328                // Check if name is provided as keyword argument
329                if let Some(name) = kwargs.swap_remove("name") {
330                    (name, vec![])
331                } else {
332                    return Err(
333                        vm.new_type_error("TypeVar() missing required argument: 'name' (pos 1)")
334                    );
335                }
336            } else if args.args.len() == 1 {
337                (args.args[0].clone(), vec![])
338            } else {
339                let name = args.args[0].clone();
340                let constraints = args.args[1..].to_vec();
341                (name, constraints)
342            };
343
344            let bound = kwargs.swap_remove("bound");
345            let covariant = kwargs
346                .swap_remove("covariant")
347                .map(|v| v.try_to_bool(vm))
348                .transpose()?
349                .unwrap_or(false);
350            let contravariant = kwargs
351                .swap_remove("contravariant")
352                .map(|v| v.try_to_bool(vm))
353                .transpose()?
354                .unwrap_or(false);
355            let infer_variance = kwargs
356                .swap_remove("infer_variance")
357                .map(|v| v.try_to_bool(vm))
358                .transpose()?
359                .unwrap_or(false);
360            let default = kwargs.swap_remove("default");
361
362            // Check for unexpected keyword arguments
363            if !kwargs.is_empty() {
364                let unexpected_keys: Vec<String> = kwargs.keys().map(|s| s.to_string()).collect();
365                return Err(vm.new_type_error(format!(
366                    "TypeVar() got unexpected keyword argument(s): {}",
367                    unexpected_keys.join(", ")
368                )));
369            }
370
371            // Check for invalid combinations
372            if covariant && contravariant {
373                return Err(vm.new_value_error("Bivariant type variables are not supported."));
374            }
375
376            if infer_variance && (covariant || contravariant) {
377                return Err(vm.new_value_error("Variance cannot be specified with infer_variance"));
378            }
379
380            // Handle constraints and bound
381            let (constraints_obj, evaluate_constraints) = if !constraints.is_empty() {
382                // Check for single constraint
383                if constraints.len() == 1 {
384                    return Err(vm.new_type_error("A single constraint is not allowed"));
385                }
386                if bound.is_some() {
387                    return Err(vm.new_type_error("Constraints cannot be used with bound"));
388                }
389                let constraints_tuple = vm.ctx.new_tuple(constraints);
390                (constraints_tuple.into(), vm.ctx.none())
391            } else {
392                (vm.ctx.none(), vm.ctx.none())
393            };
394
395            // Handle bound
396            let (bound_obj, evaluate_bound) = if let Some(bound) = bound {
397                if vm.is_none(&bound) {
398                    (vm.ctx.none(), vm.ctx.none())
399                } else {
400                    // Type check the bound
401                    let bound = type_check(bound, "Bound must be a type.", vm)?;
402                    (bound, vm.ctx.none())
403                }
404            } else {
405                (vm.ctx.none(), vm.ctx.none())
406            };
407
408            // Handle default value
409            let (default_value, evaluate_default) = if let Some(default) = default {
410                (default, vm.ctx.none())
411            } else {
412                // If no default provided, use NoDefault singleton
413                (vm.ctx.typing_no_default.clone().into(), vm.ctx.none())
414            };
415
416            Ok(Self {
417                name,
418                bound: PyMutex::new(bound_obj),
419                evaluate_bound,
420                constraints: PyMutex::new(constraints_obj),
421                evaluate_constraints,
422                default_value: PyMutex::new(default_value),
423                evaluate_default: PyMutex::new(evaluate_default),
424                covariant,
425                contravariant,
426                infer_variance,
427            })
428        }
429    }
430
431    impl TypeVar {
432        pub fn new(
433            vm: &VirtualMachine,
434            name: PyObjectRef,
435            evaluate_bound: PyObjectRef,
436            evaluate_constraints: PyObjectRef,
437        ) -> Self {
438            Self {
439                name,
440                bound: PyMutex::new(vm.ctx.none()),
441                evaluate_bound,
442                constraints: PyMutex::new(vm.ctx.none()),
443                evaluate_constraints,
444                default_value: PyMutex::new(vm.ctx.typing_no_default.clone().into()),
445                evaluate_default: PyMutex::new(vm.ctx.none()),
446                covariant: false,
447                contravariant: false,
448                infer_variance: true,
449            }
450        }
451    }
452
453    #[pyattr]
454    #[pyclass(name = "ParamSpec", module = "typing")]
455    #[derive(Debug, PyPayload)]
456    #[allow(dead_code)]
457    pub struct ParamSpec {
458        name: PyObjectRef,
459        bound: Option<PyObjectRef>,
460        default_value: PyMutex<PyObjectRef>,
461        evaluate_default: PyMutex<PyObjectRef>,
462        covariant: bool,
463        contravariant: bool,
464        infer_variance: bool,
465    }
466
467    #[pyclass(
468        flags(HAS_DICT, HAS_WEAKREF),
469        with(AsNumber, Constructor, Representable)
470    )]
471    impl ParamSpec {
472        #[pymethod]
473        fn __mro_entries__(&self, _bases: PyObjectRef, vm: &VirtualMachine) -> PyResult {
474            Err(vm.new_type_error("Cannot subclass an instance of ParamSpec"))
475        }
476
477        #[pygetset]
478        fn __name__(&self) -> PyObjectRef {
479            self.name.clone()
480        }
481
482        #[pygetset]
483        fn args(zelf: crate::PyRef<Self>, vm: &VirtualMachine) -> PyResult {
484            let self_obj: PyObjectRef = zelf.into();
485            let psa = ParamSpecArgs {
486                __origin__: self_obj,
487            };
488            Ok(psa.into_ref(&vm.ctx).into())
489        }
490
491        #[pygetset]
492        fn kwargs(zelf: crate::PyRef<Self>, vm: &VirtualMachine) -> PyResult {
493            let self_obj: PyObjectRef = zelf.into();
494            let psk = ParamSpecKwargs {
495                __origin__: self_obj,
496            };
497            Ok(psk.into_ref(&vm.ctx).into())
498        }
499
500        #[pygetset]
501        fn __bound__(&self, vm: &VirtualMachine) -> PyObjectRef {
502            if let Some(bound) = self.bound.clone() {
503                return bound;
504            }
505            vm.ctx.none()
506        }
507
508        #[pygetset]
509        const fn __covariant__(&self) -> bool {
510            self.covariant
511        }
512
513        #[pygetset]
514        const fn __contravariant__(&self) -> bool {
515            self.contravariant
516        }
517
518        #[pygetset]
519        const fn __infer_variance__(&self) -> bool {
520            self.infer_variance
521        }
522
523        #[pygetset]
524        fn __default__(&self, vm: &VirtualMachine) -> PyResult {
525            {
526                let default_value = self.default_value.lock();
527                if !default_value.is(&vm.ctx.typing_no_default) {
528                    return Ok(default_value.clone());
529                }
530            }
531            let evaluator = self.evaluate_default.lock().clone();
532            if !vm.is_none(&evaluator) {
533                let result = evaluator.call((1i32,), vm)?;
534                *self.default_value.lock() = result.clone();
535                Ok(result)
536            } else {
537                Ok(vm.ctx.typing_no_default.clone().into())
538            }
539        }
540
541        #[pygetset]
542        fn evaluate_default(&self, vm: &VirtualMachine) -> PyResult {
543            let evaluator = self.evaluate_default.lock().clone();
544            if !vm.is_none(&evaluator) {
545                return Ok(evaluator);
546            }
547            let default_value = self.default_value.lock().clone();
548            if !default_value.is(&vm.ctx.typing_no_default) {
549                return Ok(const_evaluator_alloc(default_value, vm));
550            }
551            Ok(vm.ctx.none())
552        }
553
554        #[pymethod]
555        fn __reduce__(&self) -> PyResult {
556            Ok(self.name.clone())
557        }
558
559        #[pymethod]
560        fn has_default(&self, vm: &VirtualMachine) -> bool {
561            if !vm.is_none(&self.evaluate_default.lock()) {
562                return true;
563            }
564            !self.default_value.lock().is(&vm.ctx.typing_no_default)
565        }
566
567        #[pymethod]
568        fn __typing_subst__(
569            zelf: crate::PyRef<Self>,
570            arg: PyObjectRef,
571            vm: &VirtualMachine,
572        ) -> PyResult {
573            let self_obj: PyObjectRef = zelf.into();
574            call_typing_func_object(vm, "_paramspec_subst", (self_obj, arg))
575        }
576
577        #[pymethod]
578        fn __typing_prepare_subst__(
579            zelf: crate::PyRef<Self>,
580            alias: PyObjectRef,
581            args: PyObjectRef,
582            vm: &VirtualMachine,
583        ) -> PyResult {
584            let self_obj: PyObjectRef = zelf.into();
585            call_typing_func_object(vm, "_paramspec_prepare_subst", (self_obj, alias, args))
586        }
587    }
588
589    impl AsNumber for ParamSpec {
590        fn as_number() -> &'static PyNumberMethods {
591            static AS_NUMBER: PyNumberMethods = PyNumberMethods {
592                or: Some(|a, b, vm| {
593                    let args = PyTuple::new_ref(vec![a.to_owned(), b.to_owned()], &vm.ctx);
594                    make_union(&args, vm)
595                }),
596                ..PyNumberMethods::NOT_IMPLEMENTED
597            };
598            &AS_NUMBER
599        }
600    }
601
602    impl Constructor for ParamSpec {
603        type Args = FuncArgs;
604
605        fn slot_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
606            let mut kwargs = args.kwargs;
607            // Parse arguments manually
608            let name = if args.args.is_empty() {
609                // Check if name is provided as keyword argument
610                if let Some(name) = kwargs.swap_remove("name") {
611                    name
612                } else {
613                    return Err(
614                        vm.new_type_error("ParamSpec() missing required argument: 'name' (pos 1)")
615                    );
616                }
617            } else if args.args.len() == 1 {
618                args.args[0].clone()
619            } else {
620                return Err(vm.new_type_error("ParamSpec() takes at most 1 positional argument"));
621            };
622
623            let bound = kwargs
624                .swap_remove("bound")
625                .map(|b| type_check(b, "Bound must be a type.", vm))
626                .transpose()?;
627            let covariant = kwargs
628                .swap_remove("covariant")
629                .map(|v| v.try_to_bool(vm))
630                .transpose()?
631                .unwrap_or(false);
632            let contravariant = kwargs
633                .swap_remove("contravariant")
634                .map(|v| v.try_to_bool(vm))
635                .transpose()?
636                .unwrap_or(false);
637            let infer_variance = kwargs
638                .swap_remove("infer_variance")
639                .map(|v| v.try_to_bool(vm))
640                .transpose()?
641                .unwrap_or(false);
642            let default = kwargs.swap_remove("default");
643
644            // Check for unexpected keyword arguments
645            if !kwargs.is_empty() {
646                let unexpected_keys: Vec<String> = kwargs.keys().map(|s| s.to_string()).collect();
647                return Err(vm.new_type_error(format!(
648                    "ParamSpec() got unexpected keyword argument(s): {}",
649                    unexpected_keys.join(", ")
650                )));
651            }
652
653            // Check for invalid combinations
654            if covariant && contravariant {
655                return Err(vm.new_value_error("Bivariant type variables are not supported."));
656            }
657
658            if infer_variance && (covariant || contravariant) {
659                return Err(vm.new_value_error("Variance cannot be specified with infer_variance"));
660            }
661
662            // Handle default value
663            let default_value = default.unwrap_or_else(|| vm.ctx.typing_no_default.clone().into());
664
665            let paramspec = Self {
666                name,
667                bound,
668                default_value: PyMutex::new(default_value),
669                evaluate_default: PyMutex::new(vm.ctx.none()),
670                covariant,
671                contravariant,
672                infer_variance,
673            };
674
675            let obj = paramspec.into_ref_with_type(vm, cls)?;
676            let obj_ref: PyObjectRef = obj.into();
677            set_module_from_caller(&obj_ref, vm)?;
678            Ok(obj_ref)
679        }
680
681        fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
682            unimplemented!("use slot_new")
683        }
684    }
685
686    impl Representable for ParamSpec {
687        #[inline(always)]
688        fn repr_str(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
689            let name = zelf.__name__().str_utf8(vm)?;
690            Ok(variance_repr(
691                name.as_str(),
692                zelf.infer_variance,
693                zelf.covariant,
694                zelf.contravariant,
695            ))
696        }
697    }
698
699    impl ParamSpec {
700        pub fn new(name: PyObjectRef, vm: &VirtualMachine) -> Self {
701            Self {
702                name,
703                bound: None,
704                default_value: PyMutex::new(vm.ctx.typing_no_default.clone().into()),
705                evaluate_default: PyMutex::new(vm.ctx.none()),
706                covariant: false,
707                contravariant: false,
708                infer_variance: true,
709            }
710        }
711    }
712
713    #[pyattr]
714    #[pyclass(name = "TypeVarTuple", module = "typing")]
715    #[derive(Debug, PyPayload)]
716    #[allow(dead_code)]
717    pub struct TypeVarTuple {
718        name: PyObjectRef,
719        default_value: PyMutex<PyObjectRef>,
720        evaluate_default: PyMutex<PyObjectRef>,
721    }
722    #[pyclass(
723        flags(HAS_DICT, HAS_WEAKREF),
724        with(Constructor, Representable, Iterable)
725    )]
726    impl TypeVarTuple {
727        #[pygetset]
728        fn __name__(&self) -> PyObjectRef {
729            self.name.clone()
730        }
731
732        #[pygetset]
733        fn __default__(&self, vm: &VirtualMachine) -> PyResult {
734            {
735                let default_value = self.default_value.lock();
736                if !default_value.is(&vm.ctx.typing_no_default) {
737                    return Ok(default_value.clone());
738                }
739            }
740            let evaluator = self.evaluate_default.lock().clone();
741            if !vm.is_none(&evaluator) {
742                let result = evaluator.call((1i32,), vm)?;
743                *self.default_value.lock() = result.clone();
744                Ok(result)
745            } else {
746                Ok(vm.ctx.typing_no_default.clone().into())
747            }
748        }
749
750        #[pygetset]
751        fn evaluate_default(&self, vm: &VirtualMachine) -> PyResult {
752            let evaluator = self.evaluate_default.lock().clone();
753            if !vm.is_none(&evaluator) {
754                return Ok(evaluator);
755            }
756            let default_value = self.default_value.lock().clone();
757            if !default_value.is(&vm.ctx.typing_no_default) {
758                return Ok(const_evaluator_alloc(default_value, vm));
759            }
760            Ok(vm.ctx.none())
761        }
762
763        #[pymethod]
764        fn has_default(&self, vm: &VirtualMachine) -> bool {
765            if !vm.is_none(&self.evaluate_default.lock()) {
766                return true;
767            }
768            let default_value = self.default_value.lock();
769            !default_value.is(&vm.ctx.typing_no_default)
770        }
771
772        #[pymethod]
773        fn __reduce__(&self) -> PyObjectRef {
774            self.name.clone()
775        }
776
777        #[pymethod]
778        fn __mro_entries__(&self, _bases: PyObjectRef, vm: &VirtualMachine) -> PyResult {
779            Err(vm.new_type_error("Cannot subclass an instance of TypeVarTuple"))
780        }
781
782        #[pymethod]
783        fn __typing_subst__(&self, _arg: PyObjectRef, vm: &VirtualMachine) -> PyResult {
784            Err(vm.new_type_error("Substitution of bare TypeVarTuple is not supported"))
785        }
786
787        #[pymethod]
788        fn __typing_prepare_subst__(
789            zelf: crate::PyRef<Self>,
790            alias: PyObjectRef,
791            args: PyObjectRef,
792            vm: &VirtualMachine,
793        ) -> PyResult {
794            let self_obj: PyObjectRef = zelf.into();
795            call_typing_func_object(vm, "_typevartuple_prepare_subst", (self_obj, alias, args))
796        }
797    }
798
799    impl Iterable for TypeVarTuple {
800        fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
801            // When unpacking TypeVarTuple with *, return [Unpack[self]]
802            // This is how CPython handles Generic[*Ts]
803            let typing = vm.import("typing", 0)?;
804            let unpack = typing.get_attr("Unpack", vm)?;
805            let zelf_obj: PyObjectRef = zelf.into();
806            let unpacked = vm.call_method(&unpack, "__getitem__", (zelf_obj,))?;
807            let list = vm.ctx.new_list(vec![unpacked]);
808            let list_obj: PyObjectRef = list.into();
809            vm.call_method(&list_obj, "__iter__", ())
810        }
811    }
812
813    impl Constructor for TypeVarTuple {
814        type Args = FuncArgs;
815
816        fn slot_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
817            let mut kwargs = args.kwargs;
818            // Parse arguments manually
819            let name = if args.args.is_empty() {
820                // Check if name is provided as keyword argument
821                if let Some(name) = kwargs.swap_remove("name") {
822                    name
823                } else {
824                    return Err(vm.new_type_error(
825                        "TypeVarTuple() missing required argument: 'name' (pos 1)",
826                    ));
827                }
828            } else if args.args.len() == 1 {
829                args.args[0].clone()
830            } else {
831                return Err(vm.new_type_error("TypeVarTuple() takes at most 1 positional argument"));
832            };
833
834            let default = kwargs.swap_remove("default");
835
836            // Check for unexpected keyword arguments
837            if !kwargs.is_empty() {
838                let unexpected_keys: Vec<String> = kwargs.keys().map(|s| s.to_string()).collect();
839                return Err(vm.new_type_error(format!(
840                    "TypeVarTuple() got unexpected keyword argument(s): {}",
841                    unexpected_keys.join(", ")
842                )));
843            }
844
845            // Handle default value
846            let (default_value, evaluate_default) = if let Some(default) = default {
847                (default, vm.ctx.none())
848            } else {
849                // If no default provided, use NoDefault singleton
850                (vm.ctx.typing_no_default.clone().into(), vm.ctx.none())
851            };
852
853            let typevartuple = Self {
854                name,
855                default_value: PyMutex::new(default_value),
856                evaluate_default: PyMutex::new(evaluate_default),
857            };
858
859            let obj = typevartuple.into_ref_with_type(vm, cls)?;
860            let obj_ref: PyObjectRef = obj.into();
861            set_module_from_caller(&obj_ref, vm)?;
862            Ok(obj_ref)
863        }
864
865        fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
866            unimplemented!("use slot_new")
867        }
868    }
869
870    impl Representable for TypeVarTuple {
871        #[inline(always)]
872        fn repr_str(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
873            let name = zelf.name.str(vm)?;
874            Ok(name.to_string())
875        }
876    }
877
878    impl TypeVarTuple {
879        pub fn new(name: PyObjectRef, vm: &VirtualMachine) -> Self {
880            Self {
881                name,
882                default_value: PyMutex::new(vm.ctx.typing_no_default.clone().into()),
883                evaluate_default: PyMutex::new(vm.ctx.none()),
884            }
885        }
886    }
887
888    #[pyattr]
889    #[pyclass(name = "ParamSpecArgs", module = "typing")]
890    #[derive(Debug, PyPayload)]
891    #[allow(dead_code)]
892    pub struct ParamSpecArgs {
893        __origin__: PyObjectRef,
894    }
895    #[pyclass(with(Constructor, Representable, Comparable), flags(HAS_WEAKREF))]
896    impl ParamSpecArgs {
897        #[pymethod]
898        fn __mro_entries__(&self, _bases: PyObjectRef, vm: &VirtualMachine) -> PyResult {
899            Err(vm.new_type_error("Cannot subclass an instance of ParamSpecArgs"))
900        }
901
902        #[pygetset]
903        fn __origin__(&self) -> PyObjectRef {
904            self.__origin__.clone()
905        }
906    }
907
908    impl Constructor for ParamSpecArgs {
909        type Args = (PyObjectRef,);
910
911        fn py_new(_cls: &Py<PyType>, args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
912            let origin = args.0;
913            Ok(Self { __origin__: origin })
914        }
915    }
916
917    impl Representable for ParamSpecArgs {
918        #[inline(always)]
919        fn repr_str(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
920            // Check if origin is a ParamSpec
921            if let Ok(name) = zelf.__origin__.get_attr("__name__", vm) {
922                return Ok(format!("{name}.args", name = name.str(vm)?));
923            }
924            Ok(format!("{:?}.args", zelf.__origin__))
925        }
926    }
927
928    impl Comparable for ParamSpecArgs {
929        fn cmp(
930            zelf: &crate::Py<Self>,
931            other: &PyObject,
932            op: PyComparisonOp,
933            vm: &VirtualMachine,
934        ) -> PyResult<PyComparisonValue> {
935            op.eq_only(|| {
936                if other.class().is(zelf.class())
937                    && let Some(other_args) = other.downcast_ref::<ParamSpecArgs>()
938                {
939                    let eq = zelf.__origin__.rich_compare_bool(
940                        &other_args.__origin__,
941                        PyComparisonOp::Eq,
942                        vm,
943                    )?;
944                    return Ok(PyComparisonValue::Implemented(eq));
945                }
946                Ok(PyComparisonValue::NotImplemented)
947            })
948        }
949    }
950
951    #[pyattr]
952    #[pyclass(name = "ParamSpecKwargs", module = "typing")]
953    #[derive(Debug, PyPayload)]
954    #[allow(dead_code)]
955    pub struct ParamSpecKwargs {
956        __origin__: PyObjectRef,
957    }
958    #[pyclass(with(Constructor, Representable, Comparable), flags(HAS_WEAKREF))]
959    impl ParamSpecKwargs {
960        #[pymethod]
961        fn __mro_entries__(&self, _bases: PyObjectRef, vm: &VirtualMachine) -> PyResult {
962            Err(vm.new_type_error("Cannot subclass an instance of ParamSpecKwargs"))
963        }
964
965        #[pygetset]
966        fn __origin__(&self) -> PyObjectRef {
967            self.__origin__.clone()
968        }
969    }
970
971    impl Constructor for ParamSpecKwargs {
972        type Args = (PyObjectRef,);
973
974        fn py_new(_cls: &Py<PyType>, args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
975            let origin = args.0;
976            Ok(Self { __origin__: origin })
977        }
978    }
979
980    impl Representable for ParamSpecKwargs {
981        #[inline(always)]
982        fn repr_str(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
983            // Check if origin is a ParamSpec
984            if let Ok(name) = zelf.__origin__.get_attr("__name__", vm) {
985                return Ok(format!("{name}.kwargs", name = name.str(vm)?));
986            }
987            Ok(format!("{:?}.kwargs", zelf.__origin__))
988        }
989    }
990
991    impl Comparable for ParamSpecKwargs {
992        fn cmp(
993            zelf: &crate::Py<Self>,
994            other: &PyObject,
995            op: PyComparisonOp,
996            vm: &VirtualMachine,
997        ) -> PyResult<PyComparisonValue> {
998            op.eq_only(|| {
999                if other.class().is(zelf.class())
1000                    && let Some(other_kwargs) = other.downcast_ref::<ParamSpecKwargs>()
1001                {
1002                    let eq = zelf.__origin__.rich_compare_bool(
1003                        &other_kwargs.__origin__,
1004                        PyComparisonOp::Eq,
1005                        vm,
1006                    )?;
1007                    return Ok(PyComparisonValue::Implemented(eq));
1008                }
1009                Ok(PyComparisonValue::NotImplemented)
1010            })
1011        }
1012    }
1013
1014    /// Helper function to call typing module functions with cls as first argument
1015    /// Similar to CPython's call_typing_args_kwargs
1016    fn call_typing_args_kwargs(
1017        name: &'static str,
1018        cls: PyTypeRef,
1019        args: FuncArgs,
1020        vm: &VirtualMachine,
1021    ) -> PyResult {
1022        let typing = vm.import("typing", 0)?;
1023        let func = typing.get_attr(name, vm)?;
1024
1025        // Prepare arguments: (cls, *args)
1026        let mut call_args = vec![cls.into()];
1027        call_args.extend(args.args);
1028
1029        // Call with prepared args and original kwargs
1030        let func_args = FuncArgs {
1031            args: call_args,
1032            kwargs: args.kwargs,
1033        };
1034
1035        func.call(func_args, vm)
1036    }
1037
1038    #[pyattr]
1039    #[pyclass(name = "Generic", module = "typing")]
1040    #[derive(Debug, PyPayload)]
1041    #[allow(dead_code)]
1042    pub struct Generic;
1043
1044    #[pyclass(flags(BASETYPE, HEAPTYPE))]
1045    impl Generic {
1046        #[pyattr]
1047        fn __slots__(ctx: &Context) -> PyTupleRef {
1048            ctx.empty_tuple.clone()
1049        }
1050
1051        #[pyclassmethod]
1052        fn __class_getitem__(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
1053            call_typing_args_kwargs("_generic_class_getitem", cls, args, vm)
1054        }
1055
1056        #[pyclassmethod]
1057        fn __init_subclass__(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
1058            call_typing_args_kwargs("_generic_init_subclass", cls, args, vm)
1059        }
1060    }
1061
1062    /// Sets the default value for a type parameter, equivalent to CPython's _Py_set_typeparam_default
1063    /// This is used by the CALL_INTRINSIC_2 SetTypeparamDefault instruction
1064    pub fn set_typeparam_default(
1065        type_param: PyObjectRef,
1066        evaluate_default: PyObjectRef,
1067        vm: &VirtualMachine,
1068    ) -> PyResult {
1069        // Inner function to handle common pattern of setting evaluate_default
1070        fn try_set_default<T>(
1071            obj: &PyObject,
1072            evaluate_default: &PyObject,
1073            get_field: impl FnOnce(&T) -> &PyMutex<PyObjectRef>,
1074        ) -> bool
1075        where
1076            T: PyPayload,
1077        {
1078            if let Some(typed_obj) = obj.downcast_ref::<T>() {
1079                *get_field(typed_obj).lock() = evaluate_default.to_owned();
1080                true
1081            } else {
1082                false
1083            }
1084        }
1085
1086        // Try each type parameter type
1087        if try_set_default::<TypeVar>(&type_param, &evaluate_default, |tv| &tv.evaluate_default)
1088            || try_set_default::<ParamSpec>(&type_param, &evaluate_default, |ps| {
1089                &ps.evaluate_default
1090            })
1091            || try_set_default::<TypeVarTuple>(&type_param, &evaluate_default, |tvt| {
1092                &tvt.evaluate_default
1093            })
1094        {
1095            Ok(type_param)
1096        } else {
1097            Err(vm.new_type_error(format!(
1098                "Expected a type param, got {}",
1099                type_param.class().name()
1100            )))
1101        }
1102    }
1103}
rustpython_vm/stdlib/typevar.rs

rustpython_vm/stdlib/
typevar.rs
