Repl type checking (#319)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: samuelcolvin

crates/monty-python/python/pydantic_monty/_monty.pyi

@@ -256,12 +256,28 @@ class MontyRepl:
         *,
         script_name: str = 'main.py',
         limits: ResourceLimits | None = None,
+        type_check: bool = False,
+        type_check_stubs: str | None = None,
         dataclass_registry: list[type] | None = None,
     ) -> Self:
         """
         Create an empty REPL session ready to receive snippets via `feed_run()`.
 
         No code is parsed or executed at construction time.
+
+        Arguments:
+            script_name: Name used in tracebacks and error messages
+            limits: Optional resource limits configuration
+            type_check: Whether to type-check each snippet before execution.
+                When enabled, each `feed_run`/`feed_run_async`/`feed_start` call
+                runs static type checking before executing the code, and each
+                successfully executed snippet is appended to the accumulated
+                context used for type-checking subsequent snippets.
+            type_check_stubs: Optional stub code providing type declarations for
+                variables and functions available in the REPL, e.g. input variable
+                types or external function signatures.
+            dataclass_registry: Optional list of dataclass types to register for proper
+                isinstance() support on output.
         """
 
     @property
@@ -273,6 +289,24 @@ class MontyRepl:
         Register a dataclass type for proper isinstance() support on output.
         """
 
+    def type_check(self, code: str, prefix_code: str | None = None) -> None:
+        """
+        Perform static type checking on the given code snippet.
+
+        Checks the snippet in isolation using `prefix_code` as stub context.
+        This does not use the accumulated code from previous `feed_run` calls —
+        use `prefix_code` to provide any needed declarations.
+
+        Arguments:
+            code: The code to type check
+            prefix_code: Optional code to prepend before type checking,
+                e.g. with input variable declarations or external function signatures
+
+        Raises:
+            RuntimeError: If type checking infrastructure fails
+            MontyTypingError: If type errors are found
+        """
+
     def feed_run(
         self,
         code: str,
@@ -282,16 +316,27 @@ class MontyRepl:
         print_callback: Callable[[Literal['stdout'], str], None] | None = None,
         mount: MountDir | list[MountDir] | None = None,
         os: Callable[[str, tuple[Any, ...], dict[str, Any]], Any] | None = None,
+        skip_type_check: bool = False,
     ) -> Any:
         """
         Execute one incremental snippet and return its output.
 
-        When `inputs` is provided, the key-value pairs are injected into
-        the REPL namespace before executing the snippet.
-
-        When `external_functions` is provided, external function calls and
-        name lookups are dispatched to the provided callables — matching the
-        behavior of `Monty.run(external_functions=...)`.
+        Arguments:
+            code: The Python code snippet to execute
+            inputs: Dict of input values injected into the REPL namespace
+                before executing the snippet
+            external_functions: Dict of external function callbacks. When
+                provided, external function calls and name lookups are
+                dispatched to the provided callables — matching the behavior
+                of `Monty.run(external_functions=...)`.
+            print_callback: Optional callback for print output
+            mount: Optional filesystem mount(s) to expose inside the sandbox
+            os: Optional OS access handler for filesystem operations
+            skip_type_check: When `True`, static type checking is bypassed for
+                this snippet AND the snippet is NOT appended to the accumulated
+                type-check context, so later type-checked snippets will not see
+                any names it defined. Has no effect unless `type_check=True`
+                was set on the REPL.
         """
 
     def feed_run_async(
@@ -302,6 +347,7 @@ class MontyRepl:
         external_functions: dict[str, Callable[..., Any]] | None = None,
         print_callback: Callable[[Literal['stdout'], str], None] | None = None,
         os: AbstractOS | None = None,
+        skip_type_check: bool = False,
     ) -> Coroutine[Any, Any, Any]:
         """
         Execute one incremental snippet and return its output with support for async external functions.
@@ -315,6 +361,11 @@ class MontyRepl:
             external_functions: Dict of external function callbacks (sync or async)
             print_callback: Optional callback for print output
             os: Optional OS access handler for filesystem operations
+            skip_type_check: When `True`, static type checking is bypassed for
+                this snippet AND the snippet is NOT appended to the accumulated
+                type-check context, so later type-checked snippets will not see
+                any names it defined. Has no effect unless `type_check=True`
+                was set on the REPL.
 
         Returns:
             A coroutine that resolves to the output of the snippet
@@ -329,6 +380,7 @@ class MontyRepl:
         *,
         inputs: dict[str, Any] | None = None,
         print_callback: Callable[[Literal['stdout'], str], None] | None = None,
+        skip_type_check: bool = False,
     ) -> FunctionSnapshot | NameLookupSnapshot | FutureSnapshot | MontyComplete:
         """
         Start executing an incremental snippet, yielding snapshots for external calls.
@@ -343,6 +395,17 @@ class MontyRepl:
         including support for async external functions via `FutureSnapshot`.
 
         On completion or error, the REPL state is automatically restored.
+
+        Arguments:
+            code: The Python code snippet to execute
+            inputs: Dict of input values injected into the REPL namespace
+                before executing the snippet
+            print_callback: Optional callback for print output
+            skip_type_check: When `True`, static type checking is bypassed for
+                this snippet AND the snippet is NOT appended to the accumulated
+                type-check context, so later type-checked snippets will not see
+                any names it defined. Has no effect unless `type_check=True`
+                was set on the REPL.
         """
 
     def dump(self) -> bytes:

crates/monty-python/src/async_dispatch.rs

@@ -217,7 +217,7 @@ where
                 progress_guard.disarm();
                 return Python::attach(|py| {
                     let owner = repl_owner.bind(py).get();
-                    owner.put_repl(EitherRepl::from_core(repl));
+                    owner.put_repl_after_commit(EitherRepl::from_core(repl));
                     cleanup_notifier.finish();
                     monty_to_py(py, &value, &dc_registry)
                 });
@@ -605,7 +605,7 @@ fn restore_repl<T: ResourceTracker>(
 {
     Python::attach(|py| {
         let owner = repl_owner.bind(py).get();
-        owner.put_repl(EitherRepl::from_core(repl));
+        owner.put_repl_after_rollback(EitherRepl::from_core(repl));
     });
     cleanup_notifier.finish();
 }
@@ -635,7 +635,7 @@ where
 {
     let py_err = Python::attach(|py| {
         let owner = repl_owner.bind(py).get();
-        owner.put_repl(EitherRepl::from_core(err.repl));
+        owner.put_repl_after_rollback(EitherRepl::from_core(err.repl));
         MontyError::new_err(py, err.error)
     });
     cleanup_notifier.finish();

crates/monty-python/src/monty_cls.rs

@@ -80,7 +80,7 @@ impl PyMonty {
         let input_names = list_str(inputs, "inputs")?;
 
         if type_check {
-            py_type_check(py, &code, script_name, type_check_stubs)?;
+            py_type_check(py, &code, script_name, type_check_stubs, "type_stubs.pyi")?;
         }
 
         // Create the snapshot (parses the code)
@@ -122,7 +122,7 @@ impl PyMonty {
     /// * `MontyTypingError` if type errors are found
     #[pyo3(signature = (prefix_code=None))]
     fn type_check(&self, py: Python<'_>, prefix_code: Option<&str>) -> PyResult<()> {
-        py_type_check(py, self.runner.code(), &self.script_name, prefix_code)
+        py_type_check(py, self.runner.code(), &self.script_name, prefix_code, "type_stubs.pyi")
     }
 
     /// Executes the code and returns the result.
@@ -390,8 +390,14 @@ impl PyMonty {
     }
 }
 
-fn py_type_check(py: Python<'_>, code: &str, script_name: &str, type_stubs: Option<&str>) -> PyResult<()> {
-    let type_stubs = type_stubs.map(|type_stubs| SourceFile::new(type_stubs, "type_stubs.pyi"));
+pub(crate) fn py_type_check(
+    py: Python<'_>,
+    code: &str,
+    script_name: &str,
+    type_stubs: Option<&str>,
+    stubs_name: &str,
+) -> PyResult<()> {
+    let type_stubs = type_stubs.map(|type_stubs| SourceFile::new(type_stubs, stubs_name));
 
     let opt_diagnostics =
         type_check(&SourceFile::new(code, script_name), type_stubs.as_ref()).map_err(PyRuntimeError::new_err)?;
@@ -645,7 +651,7 @@ where
 {
     match progress {
         ReplProgress::Complete { repl, value } => {
-            repl_owner.get().put_repl(EitherRepl::from_core(repl));
+            repl_owner.get().put_repl_after_commit(EitherRepl::from_core(repl));
             PyMontyComplete::create(py, &value, &dc_registry)
         }
         ReplProgress::FunctionCall(call) => {
@@ -1352,7 +1358,7 @@ impl PyNameLookupSnapshot {
     /// `ValueError` if serialization fails.
     /// `RuntimeError` if the progress has already been resumed.
     fn dump<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
-        let bytes = serialization::dump_lookup_snapshot(&self.snapshot, &self.script_name, &self.variable_name)?;
+        let bytes = serialization::dump_lookup_snapshot(py, &self.snapshot, &self.script_name, &self.variable_name)?;
         Ok(PyBytes::new(py, &bytes))
     }
 
@@ -1597,7 +1603,7 @@ impl PyFutureSnapshot {
     /// `ValueError` if serialization fails.
     /// `RuntimeError` if the progress has already been resumed.
     fn dump<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
-        let bytes = serialization::dump_future_snapshot(&self.snapshot, &self.script_name)?;
+        let bytes = serialization::dump_future_snapshot(py, &self.snapshot, &self.script_name)?;
         Ok(PyBytes::new(py, &bytes))
     }
 
@@ -1762,7 +1768,9 @@ fn restore_repl_from_repl_start_error<T: ResourceTracker>(
 where
     EitherRepl: FromCoreRepl<T>,
 {
-    repl_owner.get().put_repl(EitherRepl::from_core(err.repl));
+    repl_owner
+        .get()
+        .put_repl_after_rollback(EitherRepl::from_core(err.repl));
     MontyError::new_err(py, err.error)
 }