monty-type-checking 1.4.0

use std::thread;

use insta::assert_snapshot;
use monty_type_checking::{ImportedModule, SnippetImports, SourceFile, TypeCheckContext, TypeChecker, scan_imports};
use monty_types::{
    CHECKER_MODULE_NAMES, ModuleStub, ModuleStubError, SANDBOX_MODULE_NAMES, TypeCheckingConfig, TypeCheckingFormat,
};

/// Type-checks one snippet with a throwaway checker, rendered in `Full`.
fn check(code: &str, path: &str) -> Option<String> {
    render(code, path, TypeCheckingConfig::default())
}

/// The `Concise` rendering of one snippet — one line per diagnostic, which
/// keeps multi-diagnostic snapshots readable.
fn check_concise(code: &str, path: &str) -> Option<String> {
    render(code, path, concise())
}

/// Diagnostics borrow the checker, so every helper renders before returning.
fn render(code: &str, path: &str, config: TypeCheckingConfig) -> Option<String> {
    let mut checker = TypeChecker::default();
    checker
        .run(&SourceFile::new(code, path), None, config)
        .expect("type check should not fail internally")
        .map(|diagnostics| diagnostics.to_string())
}

fn concise() -> TypeCheckingConfig {
    TypeCheckingConfig {
        format: TypeCheckingFormat::Concise,
        color: false,
    }
}

#[test]
fn type_checking_success() {
    let code = r"
def add(x: int, y: int) -> int:
    return x + y

result = add(1, 2)
    ";

    assert!(check(code, "main.py").is_none());
}

#[test]
fn type_checking_error() {
    let code = "\
def add(x: int, y: int) -> int:
    return x + y

result = add(1, '2')
    ";

    let error_diagnostics = check(code, "main.py").expect("expected type errors");
    assert_snapshot!(error_diagnostics, @r#"
    error[invalid-argument-type]: Argument to function `add` is incorrect
     --> main.py:4:17
      |
    2 |     return x + y
    3 |
    4 | result = add(1, '2')
      |                 ^^^ Expected `int`, found `Literal["2"]`
    info: Function defined here
     --> main.py:1:5
      |
    1 | def add(x: int, y: int) -> int:
      |     ^^^         ------ Parameter declared here
    2 |     return x + y
      |
    "#);
}

#[test]
fn type_checking_error_stubs() {
    let stubs = "\
from dataclasses import dataclass

@dataclass
class User:
    name: str
    age: int
";
    let code = "\
def add(x: int, y: int) -> int:
    return x + y

result = add(1, '2')";

    let mut checker = TypeChecker::default();
    let error_diagnostics = checker
        .run(
            &SourceFile::new(code, "main.py"),
            Some(&SourceFile::new(stubs, "type_stubs.pyi")),
            TypeCheckingConfig::default(),
        )
        .unwrap()
        .expect("expected type errors")
        .to_string();
    assert_snapshot!(error_diagnostics, @r#"
    error[invalid-argument-type]: Argument to function `add` is incorrect
     --> main.py:4:17
      |
    2 |     return x + y
    3 |
    4 | result = add(1, '2')
      |                 ^^^ Expected `int`, found `Literal["2"]`
    info: Function defined here
     --> main.py:1:5
      |
    1 | def add(x: int, y: int) -> int:
      |     ^^^         ------ Parameter declared here
    2 |     return x + y
      |
    "#);
}

/// The format is chosen before the check runs, because the diagnostics borrow
/// the checker's database and so cannot be re-rendered by a caller that only
/// holds the returned string.
#[test]
fn type_checking_error_concise() {
    let code = r"
def add(x: int, y: int) -> int:
    return x + y

result = add(1, '2')
    ";

    assert_snapshot!(
        check_concise(code, "main.py").expect("expected type errors"),
        @r#"main.py:5:17: error[invalid-argument-type] Argument to function `add` is incorrect: Expected `int`, found `Literal["2"]`"#
    );
}

/// Colour is part of the same up-front choice; only `Full` and `Concise`
/// render any.
#[test]
fn type_checking_error_color() {
    let code = "x: int = 'nope'\n";

    let colored = render(
        code,
        "main.py",
        TypeCheckingConfig {
            color: true,
            ..concise()
        },
    )
    .expect("expected type errors");
    assert!(colored.starts_with('\u{1b}'), "expected ANSI escapes, got: {colored}");
}

/// Every format name round-trips through the checker; the machine-readable
/// ones are what a host parses instead of scraping `Full`.
#[test]
fn every_format_renders() {
    let code = "x: int = 'nope'\n";

    let rendered =
        |format| render(code, "main.py", TypeCheckingConfig { format, color: false }).expect("expected type errors");

    assert_snapshot!(rendered(TypeCheckingFormat::Concise), @r#"main.py:1:10: error[invalid-assignment] Object of type `Literal["nope"]` is not assignable to `int`"#);
    assert_snapshot!(rendered(TypeCheckingFormat::Pylint), @r#"main.py:1: [invalid-assignment] Object of type `Literal["nope"]` is not assignable to `int`"#);
    assert_snapshot!(rendered(TypeCheckingFormat::Azure), @r#"##vso[task.logissue type=error;sourcepath=/main.py;linenumber=1;columnnumber=10;code=invalid-assignment;]Object of type `Literal["nope"]` is not assignable to `int`"#);
    assert_snapshot!(rendered(TypeCheckingFormat::Github), @r#"::error title=monty (invalid-assignment),file=/main.py,line=1,col=10,endLine=1,endColumn=16::main.py:1:10: invalid-assignment: Object of type `Literal["nope"]` is not assignable to `int`%0A  main.py:1:4: Declared type"#);
    assert!(rendered(TypeCheckingFormat::Json).starts_with('['));
    assert!(rendered(TypeCheckingFormat::JsonLines).starts_with('{'));
    assert!(rendered(TypeCheckingFormat::Rdjson).starts_with('{'));
    assert!(rendered(TypeCheckingFormat::Gitlab).starts_with('['));
}

#[test]
fn stdlib_datetime_resolves() {
    let code = "import datetime\nprint(datetime.datetime.now())";

    let result = check(code, "main.py");
    assert!(result.is_none(), "Expected no type errors, got: {result:#?}");
}

/// Test that good_types.py type-checks without errors.
///
/// This file uses `assert_type` from typing to verify that inferred types match expected types.
#[test]
fn type_good_types() {
    let result = check(include_str!("good_types.py"), "good_types.py");
    assert!(result.is_none(), "Expected no type errors, got: {result:#?}");
}

/// Test that bad_types.py produces the expected type errors.
///
/// Snapshot stored in `tests/snapshots/main__type_bad_types.snap`.
/// Run `cargo insta review` (or `cargo insta test --accept`) to update.
#[test]
fn type_bad_types() {
    let actual = check_concise(include_str!("bad_types.py"), "bad_types.py").expect("Expected type errors");
    assert_snapshot!(actual);
}

/// Security-critical: a checker reused across runs must evaluate each run
/// against its own source. This is the session case — every feed rewrites the
/// same script path — so a stale Salsa memo would let a name defined by an
/// earlier snippet keep resolving after it was edited away.
#[test]
fn same_path_rerun_does_not_see_previous_source() {
    let mut checker = TypeChecker::default();
    let first = checker
        .run(
            &SourceFile::new("GOOD: int = 1\nresult = GOOD + 1\n", "main.py"),
            None,
            concise(),
        )
        .unwrap()
        .map(|d| d.to_string());
    assert!(first.is_none(), "first run should succeed: {first:#?}");

    // Same path, new source: `GOOD` existed only in the previous run.
    let second = checker
        .run(&SourceFile::new("x: int = GOOD\n", "main.py"), None, concise())
        .unwrap()
        .expect("second run must error — `GOOD` was only defined in the first run")
        .to_string();
    assert_snapshot!(second, @"main.py:1:10: error[unresolved-reference] Name `GOOD` used when not defined");
}

/// Security-critical: checkers hold no shared state, so concurrent sessions
/// (workers in one process, or just several threads) must never observe each
/// other's files. Guards against a process-wide cache creeping back in.
#[test]
fn concurrent_checkers_stay_isolated() {
    thread::scope(|scope| {
        let handles: Vec<_> = (0..4)
            .map(|thread_idx| {
                scope.spawn(move || {
                    let mut checker = TypeChecker::default();
                    for iter in 0..5 {
                        let defines = format!("T_{thread_idx}_{iter}: int = 1\n");
                        let ok = checker
                            .run(&SourceFile::new(&defines, "main.py"), None, concise())
                            .unwrap()
                            .map(|d| d.to_string());
                        assert!(ok.is_none(), "ok run must succeed: {ok:#?}");

                        // Same path, so the previous run's definition is gone;
                        // another thread's would only show up through shared state.
                        let probe = format!("x: int = T_{thread_idx}_{iter}\n");
                        let diagnostics = checker
                            .run(&SourceFile::new(&probe, "main.py"), None, concise())
                            .unwrap()
                            .expect("leak probe must error — prior run's name must not be visible")
                            .to_string();
                        assert_eq!(
                            diagnostics,
                            format!(
                                "main.py:1:10: error[unresolved-reference] Name `T_{thread_idx}_{iter}` used when not defined\n"
                            ),
                        );
                    }
                })
            })
            .collect();
        for handle in handles {
            handle.join().unwrap();
        }
    });
}

/// Snapshot stored in `tests/snapshots/main__reveal_types.snap`.
/// Run `cargo insta review` (or `cargo insta test --accept`) to update.
#[test]
fn test_reveal_types() {
    let actual = check_concise(include_str!("reveal_types.py"), "reveal_types.py").expect("Expected type errors");
    assert_snapshot!(actual);
}

/// ruff grows its parser stack on demand, so the type checker accepts nesting
/// the interpreter rejects; hosts reject such sources before type checking
/// (`monty::type_check_nesting_exception`) rather than rely on a limit here.
#[test]
fn deeply_nested_parentheses_type_check_without_overflow() {
    let depth = 500;
    let mut code = String::with_capacity(depth * 2 + 1);
    for _ in 0..depth {
        code.push('(');
    }
    code.push('1');
    for _ in 0..depth {
        code.push(')');
    }

    assert_eq!(check_concise(&code, "main.py"), None);
}

/// Regression test for issue #799: attribute access on a `TypedDict` value must
/// be rejected (only subscript access works at runtime). ty synthesizes
/// TypedDict members from `_typeshed._type_checker_internals.TypedDictFallback`;
/// if the vendored typeshed drops that file, attribute access silently resolves
/// to `Unknown` and this test fails with no diagnostics at all.
#[test]
fn typed_dict_attribute_access_is_rejected() {
    let stubs = "\
from typing import TypedDict

class Change(TypedDict):
    field_diffs: dict[str, object]

def get_change() -> Change: ...
";
    let code = "\
change = get_change()
ok = change['field_diffs']
bad = change.field_diffs
";
    let mut checker = TypeChecker::default();
    let diagnostics = checker
        .run(
            &SourceFile::new(code, "main.py"),
            Some(&SourceFile::new(stubs, "type_stubs.pyi")),
            concise(),
        )
        .unwrap()
        .expect("attribute access on a TypedDict must be a type error")
        .to_string();
    assert_snapshot!(diagnostics, @"main.py:3:7: error[unresolved-attribute] Object of type `Change` has no attribute `field_diffs`");
}

/// The narrowed `collections` stub must expose exactly what Monty implements at
/// runtime — importing the four implemented names type-checks clean.
#[test]
fn collections_implemented_names_type_check() {
    let code = "from collections import deque, Counter, defaultdict, namedtuple\n";
    let result = check(code, "main.py");
    assert!(result.is_none(), "expected no type errors, got: {result:?}");
}

/// Conversely, names Monty does not implement (`OrderedDict`, `ChainMap`, the
/// `User*` wrappers) are removed from the stub, so importing them is a type
/// error rather than passing the checker and then failing at runtime with
/// `ImportError`.
#[test]
fn collections_unimplemented_names_are_unresolved() {
    let code = "from collections import OrderedDict, ChainMap, UserDict, UserList, UserString\n";
    let diagnostics = check(code, "main.py").expect("expected unresolved-import errors");
    assert_snapshot!(diagnostics, @"
    error[unresolved-import]: Module `collections` has no member `OrderedDict`
     --> main.py:1:25
      |
    1 | from collections import OrderedDict, ChainMap, UserDict, UserList, UserString
      |                         ^^^^^^^^^^^

    error[unresolved-import]: Module `collections` has no member `ChainMap`
     --> main.py:1:38
      |
    1 | from collections import OrderedDict, ChainMap, UserDict, UserList, UserString
      |                                      ^^^^^^^^

    error[unresolved-import]: Module `collections` has no member `UserDict`
     --> main.py:1:48
      |
    1 | from collections import OrderedDict, ChainMap, UserDict, UserList, UserString
      |                                                ^^^^^^^^

    error[unresolved-import]: Module `collections` has no member `UserList`
     --> main.py:1:58
      |
    1 | from collections import OrderedDict, ChainMap, UserDict, UserList, UserString
      |                                                          ^^^^^^^^

    error[unresolved-import]: Module `collections` has no member `UserString`
     --> main.py:1:68
      |
    1 | from collections import OrderedDict, ChainMap, UserDict, UserList, UserString
      |                                                                    ^^^^^^^^^^
    ");
}

/// A script named like the stubs file would overwrite them and import itself, so
/// stub names vanish and spans inside the injected import underflowed (#717).
/// The db normalises paths, hence the `./` and `..` spellings.
#[test]
fn script_colliding_with_stubs_is_rejected() {
    for script_name in ["type_stubs.pyi", "./type_stubs.pyi", "sub/../type_stubs.pyi"] {
        let mut checker = TypeChecker::default();
        let error = checker
            .run(
                &SourceFile::new("from typing import TypeVar\nT = TypeVar('T\n", script_name),
                Some(&SourceFile::new("x: int\n", "type_stubs.pyi")),
                concise(),
            )
            .map(|diagnostics| diagnostics.map(|d| d.to_string()))
            .expect_err("a script colliding with the stubs file must be rejected");
        assert_eq!(
            error,
            format!("script `{script_name}` collides with the type stubs file `type_stubs.pyi`"),
        );
    }
}

// Host module stubs and the import prelude

/// The `Concise` rendering of one snippet checked with `context`.
fn render_with(checker: &mut TypeChecker, code: &str, context: &TypeCheckContext<'_>) -> Option<String> {
    checker
        .run_with(&SourceFile::new(code, "main.py"), context, concise())
        .expect("type check should not fail internally")
        .map(|diagnostics| diagnostics.to_string())
}

fn tools_stub() -> ModuleStub {
    ModuleStub::new("tools", "async def add(*, a: int, b: int) -> int: ...\n").unwrap()
}

#[test]
fn a_module_stub_resolves_the_module_import() {
    let stubs = [tools_stub()];
    let context = TypeCheckContext {
        module_stubs: &stubs,
        ..TypeCheckContext::default()
    };
    let mut checker = TypeChecker::default();
    let ok = render_with(&mut checker, "import tools\nawait tools.add(a=1, b=2)\n", &context);
    assert!(ok.is_none(), "the stubbed module should resolve: {ok:#?}");
    let wrong = render_with(&mut checker, "from tools import add\nawait add(a='x', b=2)\n", &context);
    assert_snapshot!(wrong.unwrap(), @r#"main.py:2:11: error[invalid-argument-type] Argument to function `add` is incorrect: Expected `int`, found `Literal["x"]`"#);
}

#[test]
fn a_module_stub_is_not_star_imported() {
    let stubs = [tools_stub()];
    let context = TypeCheckContext {
        module_stubs: &stubs,
        ..TypeCheckContext::default()
    };
    let mut checker = TypeChecker::default();
    let unresolved = render_with(&mut checker, "add(a=1, b=2)\n", &context);
    assert_snapshot!(unresolved.unwrap(), @"main.py:1:1: error[unresolved-reference] Name `add` used when not defined");
}

/// A `.pyi` re-exports an import only as `import x as x`, so the committed
/// snippets' star import loses `import math`; the prelude restores it.
#[test]
fn the_prelude_carries_a_committed_import_into_the_next_snippet() {
    let committed = SourceFile::new("import math\nfrom tools import add as plus\n", "repl_type_stubs.pyi");
    let stubs = [tools_stub()];
    let mut checker = TypeChecker::default();

    let without = TypeCheckContext {
        stubs: Some(&committed),
        module_stubs: &stubs,
        prelude: "",
    };
    let lost = render_with(&mut checker, "math.sqrt(4)\n", &without);
    assert_snapshot!(lost.unwrap(), @"main.py:1:1: error[unresolved-reference] Name `math` used when not defined");

    let prelude = scan_imports(committed.source_code).prelude;
    let with = TypeCheckContext {
        prelude: &prelude,
        ..without
    };
    let ok = render_with(&mut checker, "math.sqrt(4)\nawait plus(a=1, b=2)\n", &with);
    assert!(ok.is_none(), "the prelude should bind `math` and `plus`: {ok:#?}");
    // Diagnostics retain the snippet's original line numbers.
    let wrong = render_with(&mut checker, "x = 1\nmath.sqrt('4')\n", &with);
    assert_snapshot!(wrong.unwrap(), @r#"main.py:2:11: error[invalid-argument-type] Argument to function `sqrt` is incorrect: Expected `SupportsFloat | SupportsIndex`, found `Literal["4"]`"#);
}

/// Replaying imports drops `# type: ignore`; suppress the resulting diagnostics while keeping names bound.
/// Imports committed with type checking skipped also remain accepted.
#[test]
fn an_unresolved_import_in_the_prelude_is_not_reported_again() {
    let committed = SourceFile::new("import requests\n", "repl_type_stubs.pyi");
    let prelude = scan_imports("import requests  # type: ignore\n").prelude;
    assert_eq!(prelude, "import requests\n");
    let context = TypeCheckContext {
        stubs: Some(&committed),
        module_stubs: &[],
        prelude: &prelude,
    };
    let mut checker = TypeChecker::default();
    let ok = render_with(&mut checker, "x = 1\nrequests.get('/')\n", &context);
    assert!(ok.is_none(), "the prelude's own diagnostics must not surface: {ok:#?}");
    // The snippet's own unresolved imports must still be reported.
    let wrong = render_with(&mut checker, "x = 1\nimport requests\n", &context);
    assert_snapshot!(wrong.unwrap(), @"main.py:2:8: error[unresolved-import] Cannot resolve imported module `requests`");
}

/// Module-level branches may bind globals; function and class imports only bind locally.
#[test]
fn scan_imports_keeps_aliases_and_drops_the_rest() {
    let source = "\
from __future__ import annotations
import math, json as j
from tools import add, sub as minus
from . import sibling
x = 1
def f():
    import re
    return re
class C:
    import re as class_re
if x:
    import os
else:
    import sys
for _ in ():
    import time
try:
    import fast
except ImportError:
    import slow as fast
finally:
    import datetime
with open('f'):
    import random
match x:
    case 1:
        import asyncio
";
    let imports = scan_imports(source);
    assert_snapshot!(imports.prelude, @r"
    import math, json as j
    from tools import add, sub as minus
    import os
    import sys
    import time
    import fast
    import slow as fast
    import datetime
    import random
    import asyncio
    ");
    let modules: Vec<_> = imports.modules.iter().map(|module| module.name.as_str()).collect();
    assert_eq!(
        modules,
        [
            "math", "json", "tools", "os", "sys", "time", "fast", "slow", "datetime", "random", "asyncio"
        ]
    );
    assert_eq!(scan_imports("x = (\n"), SnippetImports::default());

    // Dotted imports share their top-level module's first import location.
    let imports = scan_imports("import a.b\nfrom a.c import d\nimport a\n");
    assert_eq!(imports.modules.len(), 1);
    let ImportedModule { name, range } = &imports.modules[0];
    assert_eq!(name, "a");
    assert_eq!(&"import a.b\nfrom a.c import d\nimport a\n"[*range], "import a.b");
}

#[test]
fn module_stub_names_are_validated() {
    let stub = |name: &str| ModuleStub::new(name, "").map(|_| ()).unwrap_err();
    assert_eq!(stub("1tools"), ModuleStubError::InvalidName("1tools".to_owned()));
    assert_eq!(stub("class"), ModuleStubError::InvalidName("class".to_owned()));
    assert_eq!(stub("a.1b"), ModuleStubError::InvalidName("a.1b".to_owned()));
    assert_eq!(stub("a."), ModuleStubError::InvalidName("a.".to_owned()));
    assert_eq!(stub("json"), ModuleStubError::ReservedName("json".to_owned()));
    assert_eq!(stub("json.decoder"), ModuleStubError::ReservedName("json".to_owned()));
    assert_eq!(stub("builtins"), ModuleStubError::ReservedName("builtins".to_owned()));
    assert_snapshot!(stub("a.1b").to_string(), @r#"module name "a.1b" is not a valid identifier"#);
    assert_snapshot!(stub("json").to_string(), @r#"module "json" is provided by the sandbox or its type checker and cannot be replaced"#);
    assert_eq!(ModuleStub::new("stripe", "x: int\n").unwrap().module(), "stripe");
    // Only top-level names are reserved.
    assert_eq!(ModuleStub::new("stripe.json", "").unwrap().module(), "stripe.json");
}

/// Missing ancestors become empty packages; submodule members stay separate from parent members.
#[test]
fn dotted_module_stubs_form_a_package() {
    let stubs = [
        ModuleStub::new("pkg", "VERSION: int\n").unwrap(),
        ModuleStub::new("pkg.tools", "def add(a: int, b: int) -> int: ...\n").unwrap(),
        ModuleStub::new("bare.sub", "x: str\n").unwrap(),
    ];
    let context = TypeCheckContext {
        module_stubs: &stubs,
        ..TypeCheckContext::default()
    };
    let mut checker = TypeChecker::default();
    let code = "import pkg.tools\nfrom pkg.tools import add\nfrom pkg import tools, VERSION\nimport bare.sub\n[pkg.tools.add(1, 2), add(3, 4), tools.add(5, 6), VERSION, bare.sub.x]\n";
    let ok = render_with(&mut checker, code, &context);
    assert!(ok.is_none(), "the package stubs should resolve: {ok:#?}");
    let wrong = render_with(&mut checker, "from pkg.tools import add\nadd('x', 2)\n", &context);
    assert_snapshot!(wrong.unwrap(), @r#"main.py:2:5: error[invalid-argument-type] Argument to function `add` is incorrect: Expected `int`, found `Literal["x"]`"#);
    let unresolved = render_with(&mut checker, "from pkg import add\n", &context);
    assert_snapshot!(unresolved.unwrap(), @"main.py:1:17: error[unresolved-import] Module `pkg` has no member `add`");
}

/// Host stubs must not shadow the checker's vendored typeshed definitions.
#[test]
fn reserved_names_cover_the_vendored_typeshed() {
    let versions = monty_typeshed::file_system()
        .read_to_string("stdlib/VERSIONS")
        .expect("the vendored typeshed lists its modules");
    for line in versions.lines() {
        let entry = line.split('#').next().unwrap_or_default().trim();
        let Some((module, _)) = entry.split_once(':') else {
            continue;
        };
        let top_level = module.split('.').next().unwrap_or_default();
        assert!(
            SANDBOX_MODULE_NAMES.contains(&top_level) || CHECKER_MODULE_NAMES.contains(&top_level),
            "typeshed module {top_level} is missing from SANDBOX_MODULE_NAMES and CHECKER_MODULE_NAMES"
        );
    }
    // Keep runtime and checker-only reservations separate.
    for name in CHECKER_MODULE_NAMES {
        assert!(!SANDBOX_MODULE_NAMES.contains(name), "{name} is in both lists");
    }
}

#[test]
fn bisect_star_import_exports_the_public_functions() {
    let code = "\
from bisect import *
assert bisect_left([1, 2], 2) == 1
assert bisect_right([1, 2], 2) == 2
assert bisect([1, 2], 2) == 2
insort_left([1, 2], 2)
insort_right([1, 2], 2)
insort([1, 2], 2)
";
    assert!(check(code, "main.py").is_none());
}

#[test]
fn bisect_does_not_declare_runtime_all_metadata() {
    let diagnostics = check_concise("import bisect\nbisect.__all__", "main.py").expect("__all__ is absent at runtime");
    assert_snapshot!(diagnostics, @"main.py:2:1: error[unresolved-attribute] Module `bisect` has no member `__all__`");
}