rustpython_vm/builtins/
namespace.rs

1use super::{PyStr, PyTupleRef, PyType, tuple::IntoPyTuple};
2use crate::{
3    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
4    builtins::PyDict,
5    class::PyClassImpl,
6    function::{FuncArgs, PyComparisonValue},
7    recursion::ReprGuard,
8    types::{
9        Comparable, Constructor, DefaultConstructor, Initializer, PyComparisonOp, Representable,
10    },
11};
12use rustpython_common::wtf8::Wtf8Buf;
13
14/// A simple attribute-based namespace.
15///
16/// SimpleNamespace(**kwargs)
17#[pyclass(module = "types", name = "SimpleNamespace")]
18#[derive(Debug, Default)]
19pub struct PyNamespace {}
20
21impl PyPayload for PyNamespace {
22    #[inline]
23    fn class(ctx: &Context) -> &'static Py<PyType> {
24        ctx.types.namespace_type
25    }
26}
27
28impl DefaultConstructor for PyNamespace {}
29
30#[pyclass(
31    flags(BASETYPE, HAS_DICT, HAS_WEAKREF),
32    with(Constructor, Initializer, Comparable, Representable)
33)]
34impl PyNamespace {
35    #[pymethod]
36    fn __reduce__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyTupleRef {
37        let dict = zelf.as_object().dict().unwrap();
38        let obj = zelf.as_object().to_owned();
39        let result: (PyObjectRef, PyObjectRef, PyObjectRef) = (
40            obj.class().to_owned().into(),
41            vm.new_tuple(()).into(),
42            dict.into(),
43        );
44        result.into_pytuple(vm)
45    }
46
47    #[pymethod]
48    fn __replace__(zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
49        if !args.args.is_empty() {
50            return Err(vm.new_type_error("__replace__() takes no positional arguments"));
51        }
52
53        // Create a new instance of the same type
54        let cls: PyObjectRef = zelf.class().to_owned().into();
55        let result = cls.call((), vm)?;
56
57        // Copy the current namespace dict to the new instance
58        let src_dict = zelf.dict().unwrap();
59        let dst_dict = result.dict().unwrap();
60        for (key, value) in src_dict {
61            dst_dict.set_item(&*key, value, vm)?;
62        }
63
64        // Update with the provided kwargs
65        for (name, value) in args.kwargs {
66            let name = vm.ctx.new_str(name);
67            result.set_attr(&name, value, vm)?;
68        }
69
70        Ok(result)
71    }
72}
73
74impl Initializer for PyNamespace {
75    type Args = FuncArgs;
76
77    fn init(zelf: PyRef<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
78        // SimpleNamespace accepts 0 or 1 positional argument (a mapping)
79        if args.args.len() > 1 {
80            return Err(vm.new_type_error(format!(
81                "{} expected at most 1 positional argument, got {}",
82                zelf.class().name(),
83                args.args.len()
84            )));
85        }
86
87        // If there's a positional argument, treat it as a mapping
88        if let Some(mapping) = args.args.first() {
89            // Convert to dict if not already
90            let dict: PyRef<PyDict> = if let Some(d) = mapping.downcast_ref::<PyDict>() {
91                d.to_owned()
92            } else {
93                // Call dict() on the mapping
94                let dict_type: PyObjectRef = vm.ctx.types.dict_type.to_owned().into();
95                dict_type
96                    .call((mapping.clone(),), vm)?
97                    .downcast()
98                    .map_err(|_| vm.new_type_error("dict() did not return a dict"))?
99            };
100
101            // Validate keys are strings and set attributes
102            for (key, value) in dict.into_iter() {
103                let key_str = key
104                    .downcast_ref::<crate::builtins::PyStr>()
105                    .ok_or_else(|| {
106                        vm.new_type_error(format!(
107                            "keywords must be strings, not '{}'",
108                            key.class().name()
109                        ))
110                    })?;
111                zelf.as_object().set_attr(key_str, value, vm)?;
112            }
113        }
114
115        // Apply keyword arguments (these override positional mapping values)
116        for (name, value) in args.kwargs {
117            let name = vm.ctx.new_str(name);
118            zelf.as_object().set_attr(&name, value, vm)?;
119        }
120        Ok(())
121    }
122}
123
124impl Comparable for PyNamespace {
125    fn cmp(
126        zelf: &Py<Self>,
127        other: &PyObject,
128        op: PyComparisonOp,
129        vm: &VirtualMachine,
130    ) -> PyResult<PyComparisonValue> {
131        let other = class_or_notimplemented!(Self, other);
132        let (d1, d2) = (
133            zelf.as_object().dict().unwrap(),
134            other.as_object().dict().unwrap(),
135        );
136        PyDict::cmp(&d1, d2.as_object(), op, vm)
137    }
138}
139
140impl Representable for PyNamespace {
141    #[inline]
142    fn repr_wtf8(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
143        let o = zelf.as_object();
144        let name = if o.class().is(vm.ctx.types.namespace_type) {
145            "namespace".to_owned()
146        } else {
147            o.class().slot_name().to_owned()
148        };
149
150        let repr = if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
151            let dict = zelf.as_object().dict().unwrap();
152            let mut result = Wtf8Buf::from(format!("{name}("));
153            let mut first = true;
154            for (key, value) in dict {
155                let Some(key_str) = key.downcast_ref::<PyStr>() else {
156                    continue;
157                };
158                if key_str.as_wtf8().is_empty() {
159                    continue;
160                }
161                if !first {
162                    result.push_str(", ");
163                }
164                first = false;
165                result.push_wtf8(key_str.as_wtf8());
166                result.push_char('=');
167                result.push_wtf8(value.repr(vm)?.as_wtf8());
168            }
169            result.push_char(')');
170            result
171        } else {
172            Wtf8Buf::from(format!("{name}(...)"))
173        };
174        Ok(repr)
175    }
176}
177
178pub fn init(context: &'static Context) {
179    PyNamespace::extend_class(context, context.types.namespace_type);
180}
rustpython_vm/builtins/namespace.rs

rustpython_vm/builtins/
namespace.rs
