johnandersen777.bsky.social/asyncai

Switch repository

Loading…

//! 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<String> {
    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<String> {
    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<String>, 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<String>, 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<usize> {
    let bytes = text.as_bytes();
    let mut depth = 0i32;
    let mut quote: Option<u8> = 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<char> = 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<String> {
    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"));
    }
}