//! What a snippet binds at module level: used to reject helper shadowing and //! to list the variables a REPL scope holds. //! //! Line based on purpose: only unindented statements bind module names, and //! a miss only means a variable is not listed. /// Built-in helper functions from `helpers.py`; snippets may not rebind them. pub const HELPER_NAMES: &[&str] = &["items", "names", "dig", "count_by", "to_json"]; /// Most variables listed in the scope note. const MAX_LISTED: usize = 30; /// Longest scope note. const MAX_NOTE_CHARS: usize = 600; /// Module-level names bound by `code`, in order of first appearance. pub fn assigned_names(code: &str) -> Vec { let mut names = Vec::new(); let mut in_string: Option<&str> = None; for line in code.lines() { if let Some(quote) = in_string { if line.matches(quote).count() % 2 == 1 { in_string = None; } continue; } for quote in ["\"\"\"", "'''"] { if line.matches(quote).count() % 2 == 1 { in_string = Some(quote); } } if line.starts_with(char::is_whitespace) || line.starts_with('#') { continue; } for name in line_bindings(line.trim_end()) { if !names.contains(&name) { names.push(name); } } } names } /// The first helper name `code` rebinds at module level. pub fn shadowed_helper(code: &str) -> Option<&'static str> { let names = assigned_names(code); HELPER_NAMES .iter() .copied() .find(|helper| names.iter().any(|name| name == helper)) } fn line_bindings(line: &str) -> Vec { let mut out = Vec::new(); let words: Vec<&str> = line.split_whitespace().collect(); match words.first().copied() { Some("def") | Some("class") => push_ident(&mut out, ident_prefix(words.get(1))), Some("async") if words.get(1) == Some(&"def") => { push_ident(&mut out, ident_prefix(words.get(2))) } Some("async") if words.get(1) == Some(&"for") => { if let Some(targets) = between(line, "for ", " in ") { push_targets(&mut out, targets); } } Some("for") => { if let Some(targets) = between(line, "for ", " in ") { push_targets(&mut out, targets); } } Some("import") => { for part in line["import".len()..].split(',') { let part = part.trim(); match part.split_once(" as ") { Some((_, alias)) => push_ident(&mut out, alias), None => push_ident(&mut out, part.split('.').next().unwrap_or("")), } } } Some("from") => { if let Some((_, imported)) = line.split_once(" import ") { for part in imported.trim_matches(|c| c == '(' || c == ')').split(',') { let part = part.trim(); match part.split_once(" as ") { Some((_, alias)) => push_ident(&mut out, alias), None => push_ident(&mut out, part), } } } } Some("with") | Some("async") => { let mut rest = line; while let Some((_, after)) = rest.split_once(" as ") { let target: String = after .chars() .take_while(|c| c.is_alphanumeric() || *c == '_') .collect(); push_ident(&mut out, &target); rest = after; } } Some( "if" | "elif" | "else" | "while" | "try" | "except" | "finally" | "return" | "raise" | "assert" | "del" | "pass" | "print" | "await" | "global" | "nonlocal", ) => {} _ => { let mut rest = line; while let Some(index) = assignment_index(rest) { let target = &rest[..index]; let target = match split_top(target, ':').first() { Some(before) => *before, None => target, }; push_targets(&mut out, target); rest = &rest[index + 1..]; } } } out } fn ident_prefix<'a>(word: Option<&&'a str>) -> &'a str { let word = word.copied().unwrap_or(""); let end = word .find(|c: char| !(c.is_alphanumeric() || c == '_')) .unwrap_or(word.len()); &word[..end] } fn between<'a>(line: &'a str, start: &str, end: &str) -> Option<&'a str> { let after = line.split_once(start)?.1; Some(after.split_once(end)?.0) } fn push_targets(out: &mut Vec, targets: &str) { for target in split_top(targets, ',') { let target = target .trim() .trim_start_matches('*') .trim_matches(|c| c == '(' || c == ')' || c == '[' || c == ']') .trim(); if target.contains(',') { push_targets(out, target); } else { push_ident(out, target); } } } fn push_ident(out: &mut Vec, candidate: &str) { let candidate = candidate.trim().trim_end_matches(':'); let mut chars = candidate.chars(); let valid = chars .next() .is_some_and(|first| first.is_alphabetic() || first == '_') && chars.all(|c| c.is_alphanumeric() || c == '_'); if valid && !out.iter().any(|known| known == candidate) { out.push(candidate.to_string()); } } /// Byte index of the first top-level `=` that is an assignment. fn assignment_index(text: &str) -> Option { let bytes = text.as_bytes(); let mut depth = 0i32; let mut quote: Option = None; let mut index = 0; while index < bytes.len() { let byte = bytes[index]; if let Some(open) = quote { if byte == b'\\' { index += 2; continue; } if byte == open { quote = None; } index += 1; continue; } match byte { b'"' | b'\'' => quote = Some(byte), b'#' => return None, b'(' | b'[' | b'{' => depth += 1, b')' | b']' | b'}' => depth -= 1, b'=' if depth == 0 => { let next = bytes.get(index + 1).copied(); let previous = index.checked_sub(1).map(|i| bytes[i]); let compound = previous.is_some_and(|p| b"=!<>+-*/%&|^:@".contains(&p)); if next == Some(b'=') { index += 2; continue; } if !compound { return Some(index); } return None; } _ => {} } index += 1; } None } /// Splits on `separator` outside brackets and strings. fn split_top(text: &str, separator: char) -> Vec<&str> { let mut parts = Vec::new(); let mut depth = 0i32; let mut quote: Option = None; let mut start = 0; for (index, c) in text.char_indices() { if let Some(open) = quote { if c == open { quote = None; } continue; } match c { '"' | '\'' => quote = Some(c), '(' | '[' | '{' => depth += 1, ')' | ']' | '}' => depth -= 1, _ if c == separator && depth == 0 => { parts.push(&text[start..index]); start = index + c.len_utf8(); } _ => {} } } parts.push(&text[start..]); parts } /// Python that describes each candidate name still bound in the scope. /// Its last expression is a dict `{name: "list[80]"}`. pub fn describe_program(candidates: &[String]) -> String { let mut code = String::from("_asyncai_vars = {}\n"); for name in candidates { code.push_str(&format!( "try:\n _asyncai_vars[{name:?}] = _asyncai_describe({name})\nexcept NameError:\n pass\n" )); } code.push_str("_asyncai_vars\n"); code } /// Variable names worth tracking: not private, not helpers, not pre-imports. pub fn trackable(name: &str) -> bool { !name.starts_with('_') && !HELPER_NAMES.contains(&name) && !matches!(name, "asyncio" | "json" | "re") } /// The scope note shown to the model, or `None` when nothing is bound. pub fn render_note(variables: &[(String, String)]) -> Option { if variables.is_empty() { return None; } let mut note = String::from("[run_code variables still in scope: "); for (listed, (name, description)) in variables.iter().enumerate() { let entry = format!("{name}: {description}"); if listed == MAX_LISTED || note.len() + entry.len() > MAX_NOTE_CHARS { note.push_str(&format!(", +{} more", variables.len() - listed)); break; } if listed > 0 { note.push_str(", "); } note.push_str(&entry); } note.push_str(". Reuse them instead of fetching again.]"); Some(note) } #[cfg(test)] mod tests { use super::*; #[test] fn finds_module_level_bindings() { let code = r#" import asyncio, json as j from collections import Counter, defaultdict as dd ss = await load("statesyncs") a, (b, c) = 1, (2, 3) x: int = 5 p = q = 0 total += 1 if x == 3: inner = 1 for i, item in enumerate(ss): pass async def load(res): r = 1 def helper(x=1): return x with open_thing() as handle: pass d["k"] = 1 obj.attr = 2 flag = a == b "#; assert_eq!( assigned_names(code), [ "asyncio", "j", "Counter", "dd", "ss", "a", "b", "c", "x", "p", "q", "i", "item", "load", "helper", "handle", "flag" ] ); } #[test] fn ignores_triple_quoted_text_and_keywords() { let code = "s = \"\"\"\nnot_a_var = 1\n\"\"\"\nprint(x == 1)\nreturn_value = 2"; assert_eq!(assigned_names(code), ["s", "return_value"]); } #[test] fn detects_shadowed_helpers() { assert_eq!(shadowed_helper("names = []\n"), Some("names")); assert_eq!(shadowed_helper("def dig(x):\n return x\n"), Some("dig")); assert_eq!(shadowed_helper("x = names(items)\n"), None); assert_eq!(shadowed_helper(" names = []\n"), None); } #[test] fn note_is_bounded() { assert_eq!(render_note(&[]), None); let note = render_note(&[("ss".into(), "list[80]".into())]).unwrap(); assert!(note.contains("ss: list[80]")); let many: Vec<(String, String)> = (0..100) .map(|i| (format!("variable_{i}"), "dict[3]".to_string())) .collect(); let note = render_note(&many).unwrap(); assert!(note.len() < MAX_NOTE_CHARS + 100); assert!(note.contains("more")); } }