rustpython_vm/types/

structseq.rs

use crate::common::lock::LazyLock;
use crate::{
    AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func,
    builtins::{
        PyBaseExceptionRef, PyDict, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef,
    },
    class::{PyClassImpl, StaticType, class_attr_item_doc},
    function::{
        Either, FuncArgs, KwArgs, NameChanges, OptionalArg, PyComparisonValue, PyMethodDef,
        PyMethodFlags,
    },
    iter::PyExactSizeIterator,
    protocol::{PyMappingMethods, PySequenceMethods},
    sliceable::{SequenceIndex, SliceableSequenceOp},
    types::PyComparisonOp,
    vm::Context,
};

const DEFAULT_STRUCTSEQ_REDUCE: PyMethodDef = PyMethodDef::new_const(
    "__reduce__",
    |zelf: PyRef<PyTuple>, vm: &VirtualMachine| -> PyTupleRef {
        vm.new_tuple((
            zelf.class().to_owned(),
            (vm.ctx.new_tuple(zelf.as_slice().to_vec()),),
        ))
    },
    PyMethodFlags::METHOD,
    crate::function::ItemDoc::static_text("__reduce__($self, /)\n--\n\n"),
);

/// Text signature `(iterable=(), /)` shared by every struct sequence.
pub const STRUCT_SEQUENCE_PARAMS: Option<&'static [crate::function::Param]> =
    Some(&[crate::function::Param {
        name: "iterable",
        kind: crate::function::ParamKind::PositionalOnly,
        default: Some(crate::function::DefaultRepr::Raw("()")),
    }]);

/// The arguments every struct sequence constructor takes.
#[derive(FromArgs)]
pub struct StructSequenceNewArgs {
    #[pyarg(any)]
    pub sequence: PyObjectRef,
    #[pyarg(any, optional, py_default = "{}")]
    pub dict: OptionalArg<PyObjectRef>,
}

/// Create a new struct sequence instance from a sequence.
///
/// `dict` supplies the hidden fields — the ones past `n_sequence_fields`, named
/// by `hidden_field_names` in order — that the sequence itself did not cover. It
/// may not name a field the sequence already supplied, nor one that does not
/// exist.
///
/// The class must have `n_sequence_fields` and `n_fields` attributes set
/// (done automatically by `PyStructSequence::extend_pyclass`).
pub fn struct_sequence_new(
    cls: PyTypeRef,
    args: StructSequenceNewArgs,
    hidden_field_names: &[&str],
    vm: &VirtualMachine,
) -> PyResult {
    // = structseq_new
    let StructSequenceNewArgs {
        sequence: seq,
        dict,
    } = args;

    #[cold]
    fn length_error(
        tp_name: &str,
        min_len: usize,
        max_len: usize,
        len: usize,
        vm: &VirtualMachine,
    ) -> PyBaseExceptionRef {
        if min_len == max_len {
            vm.new_type_error(format!(
                "{tp_name}() takes a {min_len}-sequence ({len}-sequence given)"
            ))
        } else if len < min_len {
            vm.new_type_error(format!(
                "{tp_name}() takes an at least {min_len}-sequence ({len}-sequence given)"
            ))
        } else {
            vm.new_type_error(format!(
                "{tp_name}() takes an at most {max_len}-sequence ({len}-sequence given)"
            ))
        }
    }

    let min_len: usize = cls
        .get_attr(identifier!(vm.ctx, n_sequence_fields))
        .ok_or_else(|| vm.new_type_error("missing n_sequence_fields attribute"))?
        .try_into_value(vm)?;
    let max_len: usize = cls
        .get_attr(identifier!(vm.ctx, n_fields))
        .ok_or_else(|| vm.new_type_error("missing n_fields attribute"))?
        .try_into_value(vm)?;

    let dict = match dict {
        OptionalArg::Missing => None,
        OptionalArg::Present(dict) => Some(dict.downcast::<PyDict>().map_err(|_| {
            vm.new_type_error(format!(
                "{}() takes a dict as second arg, if any",
                cls.slot_name()
            ))
        })?),
    };

    let seq: Vec<PyObjectRef> = seq.try_into_value(vm)?;
    let len = seq.len();

    if len < min_len || len > max_len {
        return Err(length_error(&cls.slot_name(), min_len, max_len, len, vm));
    }

    // Copy items and pad the hidden fields the sequence did not cover with None.
    let mut items = seq;
    items.resize_with(max_len, || vm.ctx.none());

    // Fill those padded slots from `dict`. Every key has to land in one of them:
    // a key naming a field the sequence already supplied, or no field at all,
    // would otherwise be silently dropped.
    if let Some(dict) = dict.filter(|dict| !dict.is_empty()) {
        let mut found = 0;
        let names = hidden_field_names.get(len - min_len..).unwrap_or(&[]);
        for (item, name) in items[len..].iter_mut().zip(names) {
            if let Some(value) = dict.get_item_opt(*name, vm)? {
                *item = value;
                found += 1;
            }
        }
        if found != dict.__len__() {
            return Err(vm.new_type_error(format!(
                "{}() got duplicate or unexpected field name(s)",
                cls.slot_name()
            )));
        }
    }

    PyTuple::new_unchecked(items.into_boxed_slice())
        .into_ref_with_type(vm, cls)
        .map(Into::into)
}

fn get_visible_len(obj: &PyObject, vm: &VirtualMachine) -> PyResult<usize> {
    obj.class()
        .get_attr(identifier!(vm.ctx, n_sequence_fields))
        .ok_or_else(|| vm.new_type_error("missing n_sequence_fields"))?
        .try_into_value(vm)
}

/// Sequence methods for struct sequences.
/// Uses n_sequence_fields to determine visible length.
static STRUCT_SEQUENCE_AS_SEQUENCE: LazyLock<PySequenceMethods> =
    LazyLock::new(|| PySequenceMethods {
        length: atomic_func!(|seq, vm| get_visible_len(seq.obj, vm)),
        concat: atomic_func!(|seq, other, vm| {
            // Convert to visible-only tuple, then use regular tuple concat
            let n_seq = get_visible_len(seq.obj, vm)?;
            let tuple = seq.obj.downcast_ref::<PyTuple>().unwrap();
            let visible: Vec<_> = tuple.as_slice().iter().take(n_seq).cloned().collect();
            let visible_tuple = PyTuple::new_ref(visible, &vm.ctx);
            // Use tuple's concat implementation
            visible_tuple
                .as_object()
                .sequence_unchecked()
                .concat(other, vm)
        }),
        repeat: atomic_func!(|seq, n, vm| {
            // Convert to visible-only tuple, then use regular tuple repeat
            let n_seq = get_visible_len(seq.obj, vm)?;
            let tuple = seq.obj.downcast_ref::<PyTuple>().unwrap();
            let visible: Vec<_> = tuple.as_slice().iter().take(n_seq).cloned().collect();
            let visible_tuple = PyTuple::new_ref(visible, &vm.ctx);
            // Use tuple's repeat implementation
            visible_tuple.as_object().sequence_unchecked().repeat(n, vm)
        }),
        item: atomic_func!(|seq, i, vm| {
            let n_seq = get_visible_len(seq.obj, vm)?;
            let tuple = seq.obj.downcast_ref::<PyTuple>().unwrap();
            let idx = if i < 0 {
                let pos_i = n_seq as isize + i;
                if pos_i < 0 {
                    return Err(vm.new_index_error("tuple index out of range"));
                }
                pos_i as usize
            } else {
                i as usize
            };
            if idx >= n_seq {
                return Err(vm.new_index_error("tuple index out of range"));
            }
            Ok(tuple.as_slice()[idx].clone())
        }),
        contains: atomic_func!(|seq, needle, vm| {
            let n_seq = get_visible_len(seq.obj, vm)?;
            let tuple = seq.obj.downcast_ref::<PyTuple>().unwrap();
            for item in tuple.as_slice().iter().take(n_seq) {
                if item.rich_compare_bool(needle, PyComparisonOp::Eq, vm)? {
                    return Ok(true);
                }
            }
            Ok(false)
        }),
        ..PySequenceMethods::NOT_IMPLEMENTED
    });

/// Mapping methods for struct sequences.
/// Handles subscript (indexing) with visible length bounds.
static STRUCT_SEQUENCE_AS_MAPPING: LazyLock<PyMappingMethods> =
    LazyLock::new(|| PyMappingMethods {
        length: atomic_func!(|mapping, vm| get_visible_len(mapping.obj, vm)),
        subscript: atomic_func!(|mapping, needle, vm| {
            let n_seq = get_visible_len(mapping.obj, vm)?;
            let tuple = mapping.obj.downcast_ref::<PyTuple>().unwrap();
            let visible_elements = &tuple.as_slice()[..n_seq];

            match SequenceIndex::try_from_borrowed_object(vm, needle, "tuple")? {
                SequenceIndex::Int(i) => visible_elements.getitem_by_index(vm, i),
                SequenceIndex::Slice(slice) => visible_elements
                    .getitem_by_slice(vm, slice)
                    .map(|x| vm.ctx.new_tuple(x).into()),
            }
        }),
        ..PyMappingMethods::NOT_IMPLEMENTED
    });

/// Trait for Data structs that back a PyStructSequence.
///
/// This trait is implemented by `#[pystruct_sequence_data]` on the Data struct.
/// It provides field information, tuple conversion, and element parsing.
pub trait PyStructSequenceData: Sized {
    /// Names of required fields (in order). Shown in repr.
    const REQUIRED_FIELD_NAMES: &'static [&'static str];

    /// Names of optional/skipped fields (in order, after required fields).
    const OPTIONAL_FIELD_NAMES: &'static [&'static str];

    /// Number of unnamed fields (visible but index-only access).
    const UNNAMED_FIELDS_LEN: usize = 0;

    /// Convert this Data struct into a PyTuple.
    fn into_tuple(self, vm: &VirtualMachine) -> PyTuple;

    /// Construct this Data struct from tuple elements.
    /// Default implementation returns an error.
    /// Override with `#[pystruct_sequence_data(try_from_object)]` to enable.
    fn try_from_elements(_elements: Vec<PyObjectRef>, vm: &VirtualMachine) -> PyResult<Self> {
        Err(vm.new_type_error("This struct sequence does not support construction from elements"))
    }
}

/// Trait for Python struct sequence types.
///
/// This trait is implemented by the `#[pystruct_sequence]` macro on the Python type struct.
/// It connects to the Data struct and provides Python-level functionality.
#[pyclass]
pub trait PyStructSequence: StaticType + PyClassImpl + Sized + 'static {
    /// The Data struct that provides field definitions.
    type Data: PyStructSequenceData;

    #[pyslot]
    fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
        struct_sequence_new(
            cls,
            args.bind_for(vm, Self::NAME)?,
            Self::Data::OPTIONAL_FIELD_NAMES,
            vm,
        )
    }

    /// Convert a Data struct into a PyStructSequence instance.
    fn from_data(data: Self::Data, vm: &VirtualMachine) -> PyTupleRef {
        let tuple =
            <Self::Data as ::rustpython_vm::types::PyStructSequenceData>::into_tuple(data, vm);
        let typ = Self::static_type();
        tuple
            .into_ref_with_type(vm, typ.to_owned())
            .expect("Every PyStructSequence must be a valid tuple. This is a RustPython bug.")
    }

    #[pyslot]
    fn slot_repr(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<PyStrRef> {
        let zelf = zelf
            .downcast_ref::<PyTuple>()
            .ok_or_else(|| vm.new_type_error("unexpected payload for __repr__"))?;

        let field_names = Self::Data::REQUIRED_FIELD_NAMES;
        let format_field = |(value, name): (&PyObject, _)| {
            let s = value.repr(vm)?;
            Ok(format!("{name}={s}"))
        };
        let (body, suffix) =
            if let Some(_guard) = rustpython_vm::recursion::ReprGuard::enter(vm, zelf.as_ref()) {
                let fields: PyResult<Vec<_>> = zelf
                    .as_slice()
                    .iter()
                    .map(|value| value.as_ref())
                    .zip(field_names.iter().copied())
                    .map(format_field)
                    .collect();
                (fields?.join(", "), "")
            } else {
                (String::new(), "...")
            };
        // Build qualified name: if MODULE_NAME is already in TP_NAME, use it directly.
        // Otherwise, check __module__ attribute (set by #[pymodule] at runtime).
        let type_name = if Self::MODULE_NAME.is_some() {
            alloc::borrow::Cow::Borrowed(Self::TP_NAME)
        } else {
            let typ = zelf.class();
            match typ.get_attr(identifier!(vm.ctx, __module__)) {
                Some(module) if module.downcastable::<PyStr>() => {
                    let module_str = module.downcast_ref::<PyStr>().unwrap();
                    alloc::borrow::Cow::Owned(format!("{}.{}", module_str.as_wtf8(), Self::NAME))
                }
                _ => alloc::borrow::Cow::Borrowed(Self::TP_NAME),
            }
        };
        let repr_str = format!("{type_name}({body}{suffix})");
        Ok(vm.ctx.new_str(repr_str))
    }

    // Return a copy of the structure with new values for the specified fields.
    #[pymethod]
    fn __replace__(
        zelf: PyRef<PyTuple>,
        changes: KwArgs<PyObjectRef, NameChanges>,
        vm: &VirtualMachine,
    ) -> PyResult {
        if Self::Data::UNNAMED_FIELDS_LEN > 0 {
            return Err(vm.new_type_error(format!(
                "__replace__() is not supported for {} because it has unnamed field(s)",
                zelf.class().slot_name()
            )));
        }

        let n_fields =
            Self::Data::REQUIRED_FIELD_NAMES.len() + Self::Data::OPTIONAL_FIELD_NAMES.len();
        let mut items: Vec<PyObjectRef> = zelf.as_slice()[..n_fields].to_vec();

        let mut kwargs = changes;

        // Replace fields from kwargs
        let all_field_names: Vec<&str> = Self::Data::REQUIRED_FIELD_NAMES
            .iter()
            .chain(Self::Data::OPTIONAL_FIELD_NAMES.iter())
            .copied()
            .collect();
        for (i, &name) in all_field_names.iter().enumerate() {
            if let Some(val) = kwargs.shift_remove(name) {
                items[i] = val;
            }
        }

        if !kwargs.is_empty() {
            let names = vm.ctx.new_list(
                kwargs
                    .keys()
                    .map(|k| vm.ctx.new_str(k.to_owned()).into())
                    .collect(),
            );
            let names_repr = names.as_object().repr(vm)?;
            return Err(vm.new_type_error(format!("Got unexpected field name(s): {names_repr}")));
        }

        PyTuple::new_unchecked(items.into_boxed_slice())
            .into_ref_with_type(vm, zelf.class().to_owned())
            .map(Into::into)
    }

    #[pymethod]
    fn __getitem__(zelf: PyRef<PyTuple>, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
        let n_seq = get_visible_len(zelf.as_ref(), vm)?;
        let visible_elements = &zelf.as_slice()[..n_seq];

        match SequenceIndex::try_from_borrowed_object(vm, &needle, "tuple")? {
            SequenceIndex::Int(i) => visible_elements.getitem_by_index(vm, i),
            SequenceIndex::Slice(slice) => visible_elements
                .getitem_by_slice(vm, slice)
                .map(|x| vm.ctx.new_tuple(x).into()),
        }
    }

    #[extend_class]
    fn extend_pyclass(ctx: &Context, class: &'static Py<PyType>) {
        // Getters for named visible fields (indices 0 to REQUIRED_FIELD_NAMES.len() - 1)
        for (i, &name) in Self::Data::REQUIRED_FIELD_NAMES.iter().enumerate() {
            class.set_attr(
                ctx.intern_str(name),
                ctx.new_readonly_tuple_member(name, class, i, class_attr_item_doc::<Self>(name))
                    .into(),
            );
        }

        // Getters for hidden/skipped fields (indices after visible fields)
        let visible_count = Self::Data::REQUIRED_FIELD_NAMES.len() + Self::Data::UNNAMED_FIELDS_LEN;
        for (i, &name) in Self::Data::OPTIONAL_FIELD_NAMES.iter().enumerate() {
            class.set_attr(
                ctx.intern_str(name),
                ctx.new_readonly_tuple_member(
                    name,
                    class,
                    visible_count + i,
                    class_attr_item_doc::<Self>(name),
                )
                .into(),
            );
        }

        class.set_attr(
            identifier!(ctx, __match_args__),
            ctx.new_tuple(
                Self::Data::REQUIRED_FIELD_NAMES
                    .iter()
                    .map(|&name| ctx.new_str(name).into())
                    .collect::<Vec<_>>(),
            )
            .into(),
        );

        // special fields:
        // n_sequence_fields = visible fields (named + unnamed)
        // n_fields = all fields (visible + hidden/skipped)
        // n_unnamed_fields
        let n_unnamed_fields = Self::Data::UNNAMED_FIELDS_LEN;
        let n_sequence_fields = Self::Data::REQUIRED_FIELD_NAMES.len() + n_unnamed_fields;
        let n_fields = n_sequence_fields + Self::Data::OPTIONAL_FIELD_NAMES.len();
        class.set_attr(
            identifier!(ctx, n_sequence_fields),
            ctx.new_int(n_sequence_fields).into(),
        );
        class.set_attr(identifier!(ctx, n_fields), ctx.new_int(n_fields).into());
        class.set_attr(
            identifier!(ctx, n_unnamed_fields),
            ctx.new_int(n_unnamed_fields).into(),
        );

        // Override as_sequence and as_mapping slots to use visible length
        class
            .slots
            .as_sequence
            .copy_from(&STRUCT_SEQUENCE_AS_SEQUENCE);
        class
            .slots
            .as_mapping
            .copy_from(&STRUCT_SEQUENCE_AS_MAPPING);

        // Override iter slot to return only visible elements
        class.slots.iter.store(Some(struct_sequence_iter));

        // Override hash slot to hash only visible elements
        class.slots.hash.store(Some(struct_sequence_hash));

        // Override richcompare slot to compare only visible elements
        class
            .slots
            .richcompare
            .store(Some(struct_sequence_richcompare));

        // Default __reduce__: only set if not already overridden by the impl's extend_class.
        // This allows struct sequences like sched_param to provide a custom __reduce__
        // (equivalent to METH_COEXIST in structseq.c).
        if !class.attributes.contains(ctx.intern_str("__reduce__")) {
            class.set_attr(
                ctx.intern_str("__reduce__"),
                DEFAULT_STRUCTSEQ_REDUCE.to_proper_method(class, ctx),
            );
        }
    }
}

/// Iterator function for struct sequences - returns only visible elements
fn struct_sequence_iter(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
    let tuple = zelf
        .downcast_ref::<PyTuple>()
        .ok_or_else(|| vm.new_type_error("expected tuple"))?;
    let n_seq = get_visible_len(&zelf, vm)?;
    let visible: Vec<_> = tuple.as_slice().iter().take(n_seq).cloned().collect();
    let visible_tuple = PyTuple::new_ref(visible, &vm.ctx);
    visible_tuple
        .as_object()
        .to_owned()
        .get_iter(vm)
        .map(Into::into)
}

/// Hash function for struct sequences - hashes only visible elements
fn struct_sequence_hash(
    zelf: &PyObject,
    vm: &VirtualMachine,
) -> PyResult<crate::common::hash::PyHash> {
    let tuple = zelf
        .downcast_ref::<PyTuple>()
        .ok_or_else(|| vm.new_type_error("expected tuple"))?;
    let n_seq = get_visible_len(zelf, vm)?;
    // Create a visible-only tuple and hash it
    let visible: Vec<_> = tuple.as_slice().iter().take(n_seq).cloned().collect();
    let visible_tuple = PyTuple::new_ref(visible, &vm.ctx);
    visible_tuple.as_object().hash(vm)
}

/// Rich comparison for struct sequences - compares only visible elements
fn struct_sequence_richcompare(
    zelf: &PyObject,
    other: &PyObject,
    op: PyComparisonOp,
    vm: &VirtualMachine,
) -> PyResult<Either<PyObjectRef, PyComparisonValue>> {
    let zelf_tuple = zelf
        .downcast_ref::<PyTuple>()
        .ok_or_else(|| vm.new_type_error("expected tuple"))?;

    // If other is not a tuple, return NotImplemented
    let Some(other_tuple) = other.downcast_ref::<PyTuple>() else {
        return Ok(Either::B(PyComparisonValue::NotImplemented));
    };

    let zelf_len = get_visible_len(zelf, vm)?;
    // For other, try to get visible len; if it fails (not a struct sequence), use full length
    let other_len = get_visible_len(other, vm).unwrap_or(other_tuple.as_slice().len());

    let zelf_visible = &zelf_tuple.as_slice()[..zelf_len];
    let other_visible = &other_tuple.as_slice()[..other_len];

    // Use the same comparison logic as regular tuples
    zelf_visible
        .iter()
        .map(|o| &**o)
        .richcompare(other_visible.iter().map(|o| &**o), op, vm)
        .map(|v| Either::B(PyComparisonValue::Implemented(v)))
}