interpretthis 0.4.0

Sandboxed Python AST interpreter for untrusted and LLM-generated code
Documentation
// Copyright 2026 Thomas Santerre and Moderately AI Inc.
//
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Emulation of Python's `enum` module.
//!
//! Supports `Enum`, `IntEnum`, `StrEnum`, and `auto()`. The enum
//! class machinery uses our regular class system: an enum class is
//! a registered ClassValue whose members are class attributes whose
//! values are the literal values. `auto()` returns sequential ints.
//!
//! Per-instance enum behaviour (.name, .value, identity comparison
//! via `is`) is not fully modelled — we treat enum members as their
//! underlying values, which works for the common patterns (storing
//! in dicts, comparing to literals) but loses identity-based
//! semantics.

use crate::{
    error::{EvalError, EvalResult, InterpreterError},
    value::{ExceptionValue, Value},
};

pub fn has_function(name: &str) -> bool {
    matches!(name, "auto" | "unique")
}

/// `@unique` — a class decorator that raises `ValueError` if two members share a
/// value (aliases), matching CPython. Returns the class unchanged otherwise.
fn enum_unique(state: &crate::state::InterpreterState, args: &[Value]) -> EvalResult {
    let Some(cls) = args.first() else {
        return Err(InterpreterError::TypeError(
            "unique() missing 1 required positional argument".into(),
        )
        .into());
    };
    let Value::Class(name) = cls else {
        return Err(InterpreterError::TypeError(format!(
            "{} is not an enum class",
            cls.type_name()
        ))
        .into());
    };
    if let Some(class) = state.classes.get(name) {
        // `ValueKey` carries `Box<Value>` (interior mutability via the shared
        // container handles) but is only ever used as an immutable projection
        // key here — same allowance the other value-dedup sites use.
        #[allow(clippy::mutable_key_type)]
        let mut seen: rustc_hash::FxHashMap<crate::value::ValueKey, String> =
            rustc_hash::FxHashMap::default();
        let mut dups: Vec<String> = Vec::new();
        for member in &class.enum_members {
            let value = match class.class_attrs.get(member) {
                Some(Value::EnumMember { value, .. }) => (**value).clone(),
                Some(v) => v.clone(),
                None => continue,
            };
            let Ok(key) = crate::eval::literals::value_to_key(&value) else { continue };
            if let Some(canonical) = seen.get(&key) {
                dups.push(format!("{member} -> {canonical}"));
            } else {
                seen.insert(key, member.clone());
            }
        }
        if !dups.is_empty() {
            return Err(EvalError::Exception(ExceptionValue::new(
                "ValueError",
                format!("duplicate values found in <enum '{name}'>: {}", dups.join(", ")),
            )));
        }
    }
    Ok(cls.clone())
}

/// The sentinel value `auto()` returns, replaced during enum class
/// construction. Modelled as an otherwise-unused `ModuleFunction` handle so it
/// is distinct from any real member value.
#[must_use]
pub fn auto_sentinel() -> Value {
    Value::ModuleFunction { module: "enum".into(), name: "__auto__".into() }
}

/// Whether `value` is the [`auto_sentinel`].
#[must_use]
pub fn is_auto_sentinel(value: &Value) -> bool {
    matches!(value, Value::ModuleFunction { module, name } if module == "enum" && name == "__auto__")
}

pub fn call(func: &str, _args: &[Value]) -> EvalResult {
    match func {
        "auto" => {
            // auto() returns a distinct sentinel that the enum class builder
            // (`wrap_enum_member`) replaces with the next sequential value:
            // highest previous value + 1 for int-valued enums, or the
            // lowercased member name for a StrEnum (CPython semantics). The
            // sentinel must be distinguishable from an explicit `= 1`.
            Ok(auto_sentinel())
        }
        _ => Err(InterpreterError::AttributeError(format!(
            "module 'enum' has no attribute '{func}'"
        ))
        .into()),
    }
}

/// Module-level constants — Enum, IntEnum, StrEnum classes. We model
/// them as Type sentinels that user code can inherit from via
/// `class Color(Enum):`. The class registration in eval_class_def
/// recognises these as valid base names but doesn't add behaviour
/// beyond the regular class system — the class attributes become the
/// enum members.
pub fn constant(name: &str) -> Option<Value> {
    match name {
        "Enum" | "IntEnum" | "StrEnum" | "Flag" | "IntFlag" => {
            Some(Value::Type(format!("enum.{name}")))
        }
        _ => None,
    }
}

/// `enum` module registration.
pub struct EnumModule;

#[async_trait::async_trait]
impl crate::eval::modules::Module for EnumModule {
    fn name(&self) -> &'static str {
        "enum"
    }
    fn constant(&self, name: &str) -> Option<Value> {
        constant(name)
    }
    fn has_function(&self, name: &str) -> bool {
        has_function(name)
    }
    async fn call(
        &self,
        state: &mut crate::state::InterpreterState,
        func: &str,
        args: &[Value],
        _kwargs: &indexmap::IndexMap<String, Value>,
        _tools: &crate::tools::Tools,
    ) -> EvalResult {
        if func == "unique" {
            return enum_unique(state, args);
        }
        call(func, args)
    }
}