Skip to main content

rustpython_vm/stdlib/
_typing.rs

1// spell-checker:ignore typevarobject funcobj typevartuples
2use crate::{
3    Context, PyResult, VirtualMachine, builtins::pystr::AsPyStr, class::PyClassImpl,
4    function::IntoFuncArgs,
5};
6
7pub use crate::stdlib::typevar::{
8    Generic, ParamSpec, ParamSpecArgs, ParamSpecKwargs, TypeVar, TypeVarTuple,
9    set_typeparam_default,
10};
11pub(crate) use decl::module_def;
12pub use decl::*;
13
14/// Initialize typing types (call extend_class)
15pub fn init(ctx: &'static Context) {
16    NoDefault::extend_class(ctx, ctx.types.typing_no_default_type);
17}
18
19pub fn call_typing_func_object<'a>(
20    vm: &VirtualMachine,
21    func_name: impl AsPyStr<'a>,
22    args: impl IntoFuncArgs,
23) -> PyResult {
24    let module = vm.import("typing", 0)?;
25    let func = module.get_attr(func_name.as_pystr(&vm.ctx), vm)?;
26    func.call(args, vm)
27}
28
29#[pymodule(name = "_typing", with(super::typevar::typevar))]
30pub(crate) mod decl {
31    use crate::class::PyClassDef;
32    use crate::common::lock::LazyLock;
33    use crate::{
34        AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
35        atomic_func,
36        builtins::{PyGenericAlias, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef, type_},
37        common::wtf8::Wtf8Buf,
38        function::FuncArgs,
39        protocol::{PyMappingMethods, PyNumberMethods},
40        types::{AsMapping, AsNumber, Callable, Constructor, Iterable, Representable},
41    };
42
43    #[pyfunction]
44    pub(crate) fn _idfunc(x: PyObjectRef) -> PyObjectRef {
45        x
46    }
47
48    #[pyfunction(name = "override")]
49    pub(crate) fn r#override(func: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
50        // Set __override__ attribute to True
51        // Skip the attribute silently if it is not writable.
52        // AttributeError happens if the object has __slots__ or a
53        // read-only property, TypeError if it's a builtin class.
54        let _ = func.set_attr("__override__", vm.ctx.true_value.clone(), vm);
55        func
56    }
57
58    #[pyclass(no_attr, name = "NoDefaultType", module = "typing")]
59    #[derive(Debug, PyPayload)]
60    pub struct NoDefault;
61
62    #[pyclass(with(Constructor, Representable), flags(IMMUTABLETYPE))]
63    impl Py<NoDefault> {
64        #[pymethod]
65        fn __reduce__(&self, _vm: &VirtualMachine) -> String {
66            "NoDefault".to_owned()
67        }
68    }
69
70    impl Constructor for NoDefault {
71        type Args = ();
72
73        fn slot_new(_cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
74            let _: () = args.bind_for(vm, Self::NAME)?;
75            Ok(vm.ctx.typing_no_default.clone().into())
76        }
77
78        fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
79            unreachable!("NoDefault is a singleton, use slot_new")
80        }
81    }
82
83    impl Representable for NoDefault {
84        #[inline(always)]
85        fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
86            Ok("typing.NoDefault".to_owned())
87        }
88    }
89
90    #[pyattr]
91    #[pyclass(name = "_ConstEvaluator", module = "_typing")]
92    #[derive(Debug, PyPayload)]
93    pub(crate) struct ConstEvaluator {
94        value: PyObjectRef,
95    }
96
97    #[pyclass(with(Constructor, Callable, Representable), flags(IMMUTABLETYPE))]
98    impl ConstEvaluator {}
99
100    impl Constructor for ConstEvaluator {
101        type Args = FuncArgs;
102
103        fn slot_new(_cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
104            Err(vm.new_type_error("cannot create '_typing._ConstEvaluator' instances"))
105        }
106
107        fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
108            unreachable!("ConstEvaluator cannot be instantiated from Python")
109        }
110    }
111
112    /// annotationlib.Format.STRING = 4
113    const ANNOTATE_FORMAT_STRING: i32 = 4;
114
115    impl Callable for ConstEvaluator {
116        type Args = FuncArgs;
117
118        fn call(zelf: &Py<Self>, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
119            let (format,): (i32,) = args.bind(vm)?;
120            let value = &zelf.value;
121            if format == ANNOTATE_FORMAT_STRING {
122                return typing_type_repr_value(value, vm);
123            }
124            Ok(value.clone())
125        }
126    }
127
128    /// String representation of a type for annotation purposes.
129    /// Equivalent of _Py_typing_type_repr.
130    fn typing_type_repr(obj: &PyObject, vm: &VirtualMachine) -> PyResult<String> {
131        // Ellipsis
132        if obj.is(&vm.ctx.ellipsis) {
133            return Ok("...".to_owned());
134        }
135        // NoneType -> "None"
136        if obj.is(&vm.ctx.types.none_type.as_object()) {
137            return Ok("None".to_owned());
138        }
139        // Generic aliases (has __origin__ and __args__) -> repr
140        let has_origin = obj.get_attr("__origin__", vm).is_ok();
141        let has_args = obj.get_attr("__args__", vm).is_ok();
142        if has_origin && has_args {
143            return Ok(obj.repr(vm)?.to_string());
144        }
145        // Has __qualname__ and __module__
146        if let Ok(qualname) = obj.get_attr("__qualname__", vm)
147            && let Ok(module) = obj.get_attr("__module__", vm)
148            && !vm.is_none(&module)
149            && let Some(module_str) = module.downcast_ref::<crate::builtins::PyStr>()
150        {
151            if module_str.as_bytes() == b"builtins" {
152                return Ok(qualname.str_utf8(vm)?.as_str().to_owned());
153            }
154            return Ok(format!(
155                "{}.{}",
156                module_str.as_wtf8(),
157                qualname.str_utf8(vm)?.as_str()
158            ));
159        }
160        // Fallback to repr
161        Ok(obj.repr(vm)?.to_string())
162    }
163
164    /// Format a value as a string for ANNOTATE_FORMAT_STRING.
165    /// Handles tuples specially by wrapping in parentheses.
166    fn typing_type_repr_value(value: &PyObject, vm: &VirtualMachine) -> PyResult {
167        if let Ok(tuple) = value.try_to_ref::<PyTuple>(vm) {
168            let mut parts = Vec::with_capacity(tuple.as_slice().len());
169            for item in tuple {
170                parts.push(typing_type_repr(item, vm)?);
171            }
172            let inner = if parts.len() == 1 {
173                format!("{},", parts[0])
174            } else {
175                parts.join(", ")
176            };
177            Ok(vm.ctx.new_str(format!("({inner})")).into())
178        } else {
179            Ok(vm.ctx.new_str(typing_type_repr(value, vm)?).into())
180        }
181    }
182
183    impl Representable for ConstEvaluator {
184        fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
185            let value_repr = zelf.value.repr(vm)?;
186            Ok(format!("<constevaluator {value_repr}>"))
187        }
188    }
189
190    pub(crate) fn const_evaluator_alloc(value: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
191        ConstEvaluator { value }.into_ref(&vm.ctx).into()
192    }
193
194    #[pyattr]
195    #[pyclass(name, module = "typing")]
196    #[derive(Debug, PyPayload)]
197    pub(crate) struct TypeAliasType {
198        #[pymember(name = "__name__")]
199        name: PyStrRef,
200        type_params: PyTupleRef,
201        compute_value: PyObjectRef,
202        cached_value: crate::common::lock::PyMutex<Option<PyObjectRef>>,
203        module: Option<PyObjectRef>,
204        is_lazy: bool,
205    }
206    impl TypeAliasType {
207        /// Create from intrinsic: compute_value is a callable that returns the value
208        pub(crate) fn new(
209            name: PyStrRef,
210            type_params: PyTupleRef,
211            compute_value: PyObjectRef,
212        ) -> Self {
213            Self {
214                name,
215                type_params,
216                compute_value,
217                cached_value: crate::common::lock::PyMutex::new(None),
218                module: None,
219                is_lazy: true,
220            }
221        }
222
223        /// Create with an eagerly evaluated value (used by constructor)
224        fn new_eager(
225            name: PyStrRef,
226            type_params: PyTupleRef,
227            value: PyObjectRef,
228            module: Option<PyObjectRef>,
229        ) -> Self {
230            Self {
231                name,
232                type_params,
233                compute_value: value.clone(),
234                cached_value: crate::common::lock::PyMutex::new(Some(value)),
235                module,
236                is_lazy: false,
237            }
238        }
239
240        fn __getitem__(zelf: &Py<Self>, args: PyObjectRef, vm: &VirtualMachine) -> PyResult {
241            if zelf.type_params.as_slice().is_empty() {
242                return Err(vm.new_type_error("Only generic type aliases are subscriptable"));
243            }
244            let args_tuple = if let Ok(tuple) = args.try_to_ref::<PyTuple>(vm) {
245                tuple.to_owned()
246            } else {
247                PyTuple::new_ref(vec![args], &vm.ctx)
248            };
249            let origin: PyObjectRef = zelf.as_object().to_owned();
250            Ok(PyGenericAlias::new(origin, args_tuple, false, vm)?.into_pyobject(vm))
251        }
252
253        /// Check type_params ordering: non-default params must precede default params.
254        /// Uses __default__ attribute to check if a type param has a default value,
255        /// comparing against typing.NoDefault sentinel (like get_type_param_default).
256        fn check_type_params(
257            type_params: &Py<PyTuple>,
258            vm: &VirtualMachine,
259        ) -> PyResult<Option<PyTupleRef>> {
260            if type_params.as_slice().is_empty() {
261                return Ok(None);
262            }
263            let no_default = &vm.ctx.typing_no_default;
264            let mut default_seen = false;
265            for param in type_params {
266                let dflt = param.get_attr("__default__", vm).map_err(|_| {
267                    vm.new_type_error(format!(
268                        "Expected a type param, got {}",
269                        param
270                            .repr(vm)
271                            .map_or_else(|_| "?".to_owned(), |s| s.to_string())
272                    ))
273                })?;
274                let is_no_default = dflt.is(no_default);
275                if is_no_default {
276                    if default_seen {
277                        return Err(vm.new_type_error(format!(
278                            "non-default type parameter '{}' follows default type parameter",
279                            param.repr(vm)?
280                        )));
281                    }
282                } else {
283                    default_seen = true;
284                }
285            }
286            Ok(Some(type_params.to_owned()))
287        }
288    }
289
290    #[pyclass(
291        with(Constructor, Representable, AsMapping, AsNumber, Iterable),
292        flags(IMMUTABLETYPE)
293    )]
294    impl Py<TypeAliasType> {
295        #[pygetset]
296        fn __value__(&self, vm: &VirtualMachine) -> PyResult {
297            let cached = self.cached_value.lock().clone();
298            if let Some(value) = cached {
299                return Ok(value);
300            }
301            // Call evaluator with format=1 (FORMAT_VALUE)
302            let value = self.compute_value.call((1i32,), vm)?;
303            *self.cached_value.lock() = Some(value.clone());
304            Ok(value)
305        }
306
307        #[pygetset]
308        fn __type_params__(&self) -> PyTupleRef {
309            self.type_params.clone()
310        }
311
312        #[pygetset]
313        fn __parameters__(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
314            // TypeVarTuples must be unpacked in __parameters__
315            unpack_typevartuples(&self.type_params, vm).map(|t| t.into())
316        }
317
318        #[pygetset]
319        fn __module__(&self, vm: &VirtualMachine) -> PyObjectRef {
320            if let Some(ref module) = self.module {
321                return module.clone();
322            }
323            // Fall back to compute_value's __module__ (like PyFunction_GetModule)
324            if let Ok(module) = self.compute_value.get_attr("__module__", vm) {
325                return module;
326            }
327            vm.ctx.none()
328        }
329
330        #[pymethod]
331        fn __reduce__(zelf: &Self, _vm: &VirtualMachine) -> PyObjectRef {
332            zelf.name.clone().into()
333        }
334
335        #[pymethod]
336        fn __typing_unpacked_tuple_args__(&self, vm: &VirtualMachine) -> PyObjectRef {
337            vm.ctx.none()
338        }
339
340        #[pygetset]
341        fn evaluate_value(&self, vm: &VirtualMachine) -> PyObjectRef {
342            if self.is_lazy {
343                return self.compute_value.clone();
344            }
345            const_evaluator_alloc(self.compute_value.clone(), vm)
346        }
347    }
348
349    impl Constructor for TypeAliasType {
350        type Args = FuncArgs;
351
352        fn py_new(_cls: &Py<PyType>, args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
353            // typealias(name, value, *, type_params=())
354            // name and value are positional-or-keyword; type_params is keyword-only.
355
356            // Reject unexpected keyword arguments.
357            for key in args.kwargs.keys() {
358                if !matches!(key.as_str(), Ok("name" | "value" | "type_params")) {
359                    return Err(
360                        vm.new_unexpected_keyword_type_error(Some("typealias"), &key.to_string())
361                    );
362                }
363            }
364
365            // Reject too many positional arguments
366            if args.args.len() > 2 {
367                return Err(vm.new_type_error(format!(
368                    "typealias() takes exactly 2 positional arguments ({} given)",
369                    args.args.len()
370                )));
371            }
372
373            // Resolve name: positional[0] or kwarg
374            let name = if !args.args.is_empty() {
375                if args.kwargs.contains_key("name") {
376                    return Err(vm.new_type_error(
377                        "argument for typealias() given by name ('name') and position (1)",
378                    ));
379                }
380                args.args[0].clone()
381            } else {
382                args.kwargs.get("name").cloned().ok_or_else(|| {
383                    vm.new_type_error("typealias() missing required argument 'name' (pos 1)")
384                })?
385            };
386
387            // Resolve value: positional[1] or kwarg
388            let value = if args.args.len() >= 2 {
389                if args.kwargs.contains_key("value") {
390                    return Err(vm.new_type_error(
391                        "argument for typealias() given by name ('value') and position (2)",
392                    ));
393                }
394                args.args[1].clone()
395            } else {
396                args.kwargs.get("value").cloned().ok_or_else(|| {
397                    vm.new_type_error("typealias() missing required argument 'value' (pos 2)")
398                })?
399            };
400
401            let name = name.downcast::<crate::builtins::PyStr>().map_err(|obj| {
402                vm.new_type_error(format!(
403                    "typealias() argument 'name' must be str, not {}",
404                    obj.class().name()
405                ))
406            })?;
407
408            let type_params = if let Some(tp) = args.kwargs.get("type_params") {
409                let tp = tp
410                    .clone()
411                    .downcast::<crate::builtins::PyTuple>()
412                    .map_err(|_| vm.new_type_error("type_params must be a tuple"))?;
413                Self::check_type_params(&tp, vm)?;
414                tp
415            } else {
416                vm.ctx.empty_tuple.clone()
417            };
418
419            // Get caller's module name from frame globals, like typevar.rs caller()
420            let module =
421                crate::frame::current_globals().and_then(|g| g.get_item("__name__", vm).ok());
422
423            Ok(Self::new_eager(name, type_params, value, module))
424        }
425    }
426
427    impl Representable for TypeAliasType {
428        fn repr_wtf8(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
429            Ok(zelf.name.as_wtf8().to_owned())
430        }
431    }
432
433    impl AsMapping for TypeAliasType {
434        fn as_mapping() -> &'static PyMappingMethods {
435            static AS_MAPPING: LazyLock<PyMappingMethods> = LazyLock::new(|| PyMappingMethods {
436                subscript: atomic_func!(|mapping, needle, vm| {
437                    let zelf = TypeAliasType::mapping_downcast(mapping);
438                    TypeAliasType::__getitem__(zelf, needle.to_owned(), vm)
439                }),
440                ..PyMappingMethods::NOT_IMPLEMENTED
441            });
442            &AS_MAPPING
443        }
444    }
445
446    impl AsNumber for TypeAliasType {
447        fn as_number() -> &'static PyNumberMethods {
448            static AS_NUMBER: PyNumberMethods = PyNumberMethods {
449                or: Some(|a, b, vm| type_::or_(a.to_owned(), b.to_owned(), vm)),
450                ..PyNumberMethods::NOT_IMPLEMENTED
451            };
452            &AS_NUMBER
453        }
454    }
455
456    impl Iterable for TypeAliasType {
457        fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
458            // Import typing.Unpack and return iter((Unpack[self],))
459            let typing = vm.import("typing", 0)?;
460            let unpack = typing.get_attr("Unpack", vm)?;
461            let zelf_obj: PyObjectRef = zelf.into();
462            let unpacked = vm.call_method(&unpack, "__getitem__", (zelf_obj,))?;
463            let tuple = PyTuple::new_ref(vec![unpacked], &vm.ctx);
464            Ok(tuple.as_object().get_iter(vm)?.into())
465        }
466    }
467
468    /// Wrap TypeVarTuples in Unpack[], matching unpack_typevartuples()
469    pub(crate) fn unpack_typevartuples(
470        type_params: &Py<PyTuple>,
471        vm: &VirtualMachine,
472    ) -> PyResult<PyTupleRef> {
473        let has_tvt = type_params
474            .as_slice()
475            .iter()
476            .any(|p| p.downcastable::<crate::stdlib::typevar::TypeVarTuple>());
477        if !has_tvt {
478            return Ok(type_params.to_owned());
479        }
480        let typing = vm.import("typing", 0)?;
481        let unpack_cls = typing.get_attr("Unpack", vm)?;
482        let new_params: Vec<PyObjectRef> = type_params
483            .as_slice()
484            .iter()
485            .map(|p| {
486                if p.downcastable::<crate::stdlib::typevar::TypeVarTuple>() {
487                    vm.call_method(&unpack_cls, "__getitem__", (p.clone(),))
488                } else {
489                    Ok(p.clone())
490                }
491            })
492            .collect::<PyResult<_>>()?;
493        Ok(PyTuple::new_ref(new_params, &vm.ctx))
494    }
495
496    #[expect(clippy::unnecessary_wraps, reason = "Needs to comply with a signature")]
497    pub(crate) fn module_exec(
498        vm: &VirtualMachine,
499        module: &Py<crate::builtins::PyModule>,
500    ) -> PyResult<()> {
501        __module_exec(vm, module);
502
503        extend_module!(vm, module, {
504            "NoDefault" => vm.ctx.typing_no_default.clone(),
505            "Union" => vm.ctx.types.union_type.to_owned(),
506        });
507
508        Ok(())
509    }
510}