use rustpython_parser::ast::{self, Expr};
use crate::{
error::{EvalError, EvalResult, InterpreterError},
eval::{eval_expr, place},
state::{InterpreterState, estimate_value_size},
tools::Tools,
value::{Value, shared_list},
};
const PROTECTED_BUILTINS: &[&str] = &[
"print",
"len",
"range",
"str",
"int",
"float",
"bool",
"type",
"isinstance",
"issubclass",
"super",
"hasattr",
"callable",
"abs",
"round",
"min",
"max",
"sum",
"all",
"any",
"sorted",
"enumerate",
"zip",
"reversed",
"chr",
"ord",
"list",
"tuple",
"dict",
"set",
"filter",
"map",
"object",
];
pub async fn eval_assign(
state: &mut InterpreterState,
node: &ast::StmtAssign,
tools: &Tools,
) -> EvalResult {
let value = eval_expr(state, &node.value, tools).await?;
if node.targets.len() == 1 {
assign_target(state, &node.targets[0], value, tools).await?;
} else {
for target in &node.targets {
assign_target(state, target, value.clone(), tools).await?;
}
}
Ok(Value::None)
}
pub async fn eval_aug_assign(
state: &mut InterpreterState,
node: &ast::StmtAugAssign,
tools: &Tools,
) -> EvalResult {
pretouch_defaultdict(state, &node.target, tools).await?;
let place = place::eval_place(state, &node.target, tools).await?.ok_or_else(|| {
EvalError::from(InterpreterError::Runtime("unsupported augmented assignment target".into()))
})?;
if !place.is_navigable() {
return Err(InterpreterError::Runtime(
"augmented assignment to a slice target is not supported".into(),
)
.into());
}
let current = {
let root = state
.variables
.get_mut(&place.root)
.ok_or_else(|| EvalError::from(InterpreterError::name_not_defined(&place.root)))?;
place::with_navigate_mut(root, &place.steps, |slot| slot.clone())?
};
let rhs = eval_expr(state, &node.value, tools).await?;
let new_value = crate::eval::op::aug_binop(state, node.op, ¤t, &rhs, tools).await?;
let delta = {
let root = state
.variables
.get_mut(&place.root)
.ok_or_else(|| EvalError::from(InterpreterError::name_not_defined(&place.root)))?;
place::with_navigate_mut(root, &place.steps, |slot| {
let delta =
place::size_delta(estimate_value_size(slot), estimate_value_size(&new_value));
*slot = new_value;
delta
})?
};
place::apply_mem_delta(state, delta)?;
Ok(Value::None)
}
pub(crate) async fn pretouch_defaultdict(
state: &mut InterpreterState,
target: &Expr,
tools: &Tools,
) -> Result<(), EvalError> {
pretouch_defaultdict_inner(state, target, tools, 0).await
}
fn pretouch_defaultdict_inner<'a>(
state: &'a mut InterpreterState,
target: &'a Expr,
tools: &'a Tools,
depth: u32,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), EvalError>> + Send + 'a>> {
Box::pin(async move {
if depth > 16 {
return Ok(());
}
let Expr::Subscript(sub) = target else { return Ok(()) };
pretouch_defaultdict_inner(state, sub.value.as_ref(), tools, depth + 1).await?;
let Expr::Name(name_node) = sub.value.as_ref() else { return Ok(()) };
let var_name = name_node.id.as_str().to_string();
if !matches!(state.variables.get(&var_name), Some(Value::DefaultDict(_))) {
return Ok(());
}
let key_val = eval_expr(state, &sub.slice, tools).await?;
let key = crate::eval::literals::value_to_key(&key_val)?;
let already_present = matches!(
state.variables.get(&var_name),
Some(Value::DefaultDict(data)) if data.items.contains_key(&key)
);
if already_present {
return Ok(());
}
let factory = match state.variables.get(&var_name) {
Some(Value::DefaultDict(data)) => data.factory.clone(),
_ => return Ok(()),
};
let synth = crate::eval::names::invoke_factory_pub(state, &factory, tools).await?;
if let Some(Value::DefaultDict(data)) = state.variables.get(&var_name).cloned() {
let mut new_data = *data;
new_data.items.insert(key, synth);
state
.set_variable(&var_name, Value::DefaultDict(Box::new(new_data)))
.map_err(EvalError::Interpreter)?;
}
Ok(())
})
}
pub async fn eval_ann_assign(
state: &mut InterpreterState,
node: &ast::StmtAnnAssign,
tools: &Tools,
) -> EvalResult {
if let Some(ref value_expr) = node.value {
let value = eval_expr(state, value_expr, tools).await?;
assign_target(state, &node.target, value, tools).await?;
}
Ok(Value::None)
}
async fn assign_unpacking(
state: &mut InterpreterState,
elts: &[Expr],
value: Value,
tools: &Tools,
) -> Result<(), EvalError> {
let items = if matches!(value, Value::Lazy { .. } | Value::Instance(_)) {
crate::eval::op::iter(state, &value, tools).await?
} else {
value_to_iterable(&value)?
};
let star_index = elts.iter().position(|e| matches!(e, Expr::Starred(_)));
if let Some(star_idx) = star_index {
if elts.iter().skip(star_idx + 1).any(|e| matches!(e, Expr::Starred(_))) {
return Err(InterpreterError::Runtime(
"multiple starred expressions in assignment".into(),
)
.into());
}
let non_star = elts.len() - 1;
if items.len() < non_star {
return Err(InterpreterError::Runtime(format!(
"not enough values to unpack (expected at least {non_star}, got {})",
items.len()
))
.into());
}
let tail_count = elts.len() - star_idx - 1;
let head = &items[..star_idx];
let star_slice = &items[star_idx..items.len() - tail_count];
let tail = &items[items.len() - tail_count..];
for (elem, val) in elts[..star_idx].iter().zip(head.iter()) {
assign_target(state, elem, val.clone(), tools).await?;
}
let Expr::Starred(starred) = &elts[star_idx] else {
unreachable!("star_index was set by matches! on Expr::Starred above");
};
assign_target(state, &starred.value, Value::List(shared_list(star_slice.to_vec())), tools)
.await?;
for (elem, val) in elts[star_idx + 1..].iter().zip(tail.iter()) {
assign_target(state, elem, val.clone(), tools).await?;
}
return Ok(());
}
if elts.len() != items.len() {
return Err(InterpreterError::Runtime(format!(
"not enough values to unpack (expected {}, got {})",
elts.len(),
items.len()
))
.into());
}
for (elem, val) in elts.iter().zip(items) {
assign_target(state, elem, val, tools).await?;
}
Ok(())
}
pub fn assign_target<'a>(
state: &'a mut InterpreterState,
target: &'a Expr,
value: Value,
tools: &'a Tools,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), EvalError>> + Send + 'a>> {
Box::pin(async move {
match target {
Expr::Name(name_node) => {
let name = name_node.id.as_str();
crate::security::validator::validate_name(
crate::security::validator::NameContext::Assignment,
name,
)?;
if tools.contains_key(name) {
return Err(InterpreterError::Security(format!(
"cannot assign to name '{name}': doing this would erase the existing tool"
))
.into());
}
if PROTECTED_BUILTINS.contains(&name) {
return Err(InterpreterError::Security(format!(
"cannot assign to name '{name}': doing this would erase the existing tool"
))
.into());
}
state.set_variable(name, value).map_err(EvalError::Interpreter)?;
Ok(())
}
Expr::Tuple(tuple_node) => {
assign_unpacking(state, &tuple_node.elts, value, tools).await
}
Expr::List(list_node) => assign_unpacking(state, &list_node.elts, value, tools).await,
Expr::Attribute(attr_node) => {
if let Expr::Name(name_node) = attr_node.value.as_ref() {
let obj_name = name_node.id.as_str().to_string();
if let Some(obj) = state.variables.get(&obj_name).cloned() {
if let Value::Instance(inst) = &obj {
let class_name = inst.class_name.clone();
let attr_name = attr_node.attr.as_str().to_string();
if let Some(prop) = crate::eval::classes::lookup_property(
state,
&class_name,
&attr_name,
) {
if let Some(setter) = prop.setter {
let configured = crate::eval::classes::invoke_property_setter(
state, &setter, obj, value, tools,
)
.await?;
state
.set_variable(&obj_name, configured)
.map_err(EvalError::Interpreter)?;
return Ok(());
}
if prop.setter.is_none() {
return Err(InterpreterError::AttributeError(format!(
"property '{attr_name}' of '{class_name}' object has no setter"
))
.into());
}
}
if let Some((_, method)) = crate::eval::classes::lookup_method_in_mro(
state,
&class_name,
"__setattr__",
) {
let name_arg = Value::String(attr_name.into());
let call = crate::eval::functions::CallArgs {
positional: &[name_arg, value.clone()],
keyword: &indexmap::IndexMap::new(),
};
let (_returned, updated_self) = crate::eval::classes::call_method(
state, &method, obj, call, tools,
)
.await?;
state
.set_variable(&obj_name, updated_self)
.map_err(EvalError::Interpreter)?;
return Ok(());
}
}
}
}
let target_place =
place::eval_place(state, target, tools).await?.ok_or_else(|| {
EvalError::from(InterpreterError::Runtime(
"cannot assign to this expression".into(),
))
})?;
let Some((last, prefix)) = target_place.steps.split_last() else {
return Err(InterpreterError::Runtime(
"assignment target resolved to a bare name unexpectedly".into(),
)
.into());
};
let delta = {
let root = state.variables.get_mut(&target_place.root).ok_or_else(|| {
EvalError::from(InterpreterError::name_not_defined(&target_place.root))
})?;
place::with_navigate_mut(root, prefix, |container| {
place::assign_terminal(container, last, value)
})??
};
place::apply_mem_delta(state, delta)?;
Ok(())
}
Expr::Subscript(sub) => {
let receiver = crate::eval::eval_expr(state, &sub.value, tools).await?;
if matches!(receiver, Value::Instance(_)) {
let key = crate::eval::eval_expr(state, &sub.slice, tools).await?;
if let Some(updated_self) =
crate::eval::op::setitem(state, &receiver, &key, value.clone(), tools)
.await?
{
if let Expr::Name(name_node) = sub.value.as_ref() {
state
.set_variable(name_node.id.as_str(), updated_self)
.map_err(EvalError::Interpreter)?;
}
return Ok(());
}
}
let target_place =
place::eval_place(state, target, tools).await?.ok_or_else(|| {
EvalError::from(InterpreterError::Runtime(
"cannot assign to this expression".into(),
))
})?;
let Some((last, prefix)) = target_place.steps.split_last() else {
return Err(InterpreterError::Runtime(
"assignment target resolved to a bare name unexpectedly".into(),
)
.into());
};
let delta = {
let root = state.variables.get_mut(&target_place.root).ok_or_else(|| {
EvalError::from(InterpreterError::name_not_defined(&target_place.root))
})?;
place::with_navigate_mut(root, prefix, |container| {
place::assign_terminal(container, last, value)
})??
};
place::apply_mem_delta(state, delta)?;
Ok(())
}
_ => Err(InterpreterError::Runtime(format!(
"unsupported assignment target: {:?}",
std::mem::discriminant(target)
))
.into()),
}
})
}
fn value_to_iterable(val: &Value) -> Result<Vec<Value>, EvalError> {
match val {
Value::List(items) => Ok(items.lock().clone()),
Value::Tuple(items) | Value::Set(items) => Ok(items.clone()),
Value::String(s) => Ok(s.chars().map(|c| Value::String(c.to_string().into())).collect()),
Value::Range { start, stop, step } => {
let mut items = Vec::new();
let mut i = *start;
match (*step).cmp(&0) {
std::cmp::Ordering::Greater => {
while i < *stop {
items.push(Value::Int(i));
i += step;
}
}
std::cmp::Ordering::Less => {
while i > *stop {
items.push(Value::Int(i));
i += step;
}
}
std::cmp::Ordering::Equal => {}
}
Ok(items)
}
_ => Err(InterpreterError::TypeError(format!(
"cannot unpack non-iterable {} object",
val.type_name()
))
.into()),
}
}