Visitar URL original
Fix asyncio related compiler/library issues by youknowone · Pull Request #6837 · RustPython/RustPython · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions Lib/test/test_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,8 +217,6 @@ def fun():

ctx.run(fun)

# TODO: RUSTPYTHON
@unittest.expectedFailure
@isolated_context
def test_context_getset_1(self):
c = contextvars.ContextVar('c')
Expand Down Expand Up @@ -317,8 +315,6 @@ def test_context_getset_4(self):
with self.assertRaisesRegex(ValueError, 'different Context'):
c.reset(tok)

# TODO: RUSTPYTHON
@unittest.expectedFailure
@isolated_context
def test_context_getset_5(self):
c = contextvars.ContextVar('c', default=42)
Expand All @@ -332,8 +328,6 @@ def fun():
contextvars.copy_context().run(fun)
self.assertEqual(c.get(), [])

# TODO: RUSTPYTHON
@unittest.expectedFailure
def test_context_copy_1(self):
ctx1 = contextvars.Context()
c = contextvars.ContextVar('c', default=42)
Expand Down
2 changes: 0 additions & 2 deletions Lib/test/test_inspect/test_inspect.py
Original file line number Diff line number Diff line change
Expand Up @@ -2797,7 +2797,6 @@ def test_easy_debugging(self):
self.assertIn(name, repr(state))
self.assertIn(name, str(state))

@unittest.expectedFailure # TODO: RUSTPYTHON
def test_getgeneratorlocals(self):
def each(lst, a=None):
b=(1, 2, 3)
Expand Down Expand Up @@ -2985,7 +2984,6 @@ def test_easy_debugging(self):
self.assertIn(name, repr(state))
self.assertIn(name, str(state))

@unittest.expectedFailure # TODO: RUSTPYTHON
async def test_getasyncgenlocals(self):
async def each(lst, a=None):
b=(1, 2, 3)
Expand Down
1 change: 0 additions & 1 deletion Lib/test/test_ssl.py
Original file line number Diff line number Diff line change
Expand Up @@ -3525,7 +3525,6 @@ def test_starttls(self):
else:
s.close()

@unittest.expectedFailureIfWindows("TODO: RUSTPYTHON")
def test_socketserver(self):
"""Using socketserver to create and manage SSL connections."""
server = make_https_server(self, certfile=SIGNED_CERTFILE)
Expand Down
5 changes: 1 addition & 4 deletions Lib/test/test_unittest/test_async_case.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,7 @@ class MyException(Exception):


def tearDownModule():
# XXX: RUSTPYTHON; asyncio.events._set_event_loop_policy is not implemented
# asyncio.events._set_event_loop_policy(None)
pass
asyncio.events._set_event_loop_policy(None)


class TestCM:
Expand Down Expand Up @@ -52,7 +50,6 @@ def setUp(self):
# starting a new event loop
self.addCleanup(support.gc_collect)

@unittest.expectedFailure # TODO: RUSTPYTHON
def test_full_cycle(self):
expected = ['setUp',
'asyncSetUp',
Expand Down
43 changes: 23 additions & 20 deletions crates/codegen/src/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1528,27 +1528,30 @@ impl Compiler {
// Otherwise, if an exception occurs during the finally body, the stack
// will be unwound to the wrong depth and the return value will be lost.
if preserve_tos {
// Get the handler info from the saved fblock (or current handler)
// and create a new handler with stack_depth + 1
let (handler, stack_depth, preserve_lasti) =
if let Some(handler) = saved_fblock.fb_handler {
(
Some(handler),
saved_fblock.fb_stack_depth + 1, // +1 for return value
saved_fblock.fb_preserve_lasti,
)
} else {
// No handler in saved_fblock, check current handler
if let Some(current_handler) = self.current_except_handler() {
(
Some(current_handler.handler_block),
current_handler.stack_depth + 1, // +1 for return value
current_handler.preserve_lasti,
)
} else {
(None, 1, false) // No handler, but still track the return value
// Find the outer handler, skipping ALL FinallyTry blocks.
// During return unwinding, FinallyTry handlers (finally_except_block)
// should NOT catch exceptions from the finally body - those exceptions
// should propagate to the outer TryExcept handler instead.
// FinallyTry's handler is only for the try body's exceptions.
let (handler, stack_depth, preserve_lasti) = {
let code = self.code_stack.last().unwrap();
let mut found = None;
for fblock in code.fblock.iter().rev() {
// Skip FinallyTry blocks - their handlers are for exception path
if matches!(fblock.fb_type, FBlockType::FinallyTry) {
continue;
}
if let Some(handler) = fblock.fb_handler {
found = Some((
Some(handler),
fblock.fb_stack_depth + 1, // +1 for return value
fblock.fb_preserve_lasti,
));
break;
}
};
}
found.unwrap_or((None, 1, false))
};
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

self.push_fblock_with_handler(
FBlockType::PopValue,
Expand Down
10 changes: 7 additions & 3 deletions crates/stdlib/src/contextvars.rs
Original file line number Diff line number Diff line change
Expand Up @@ -168,11 +168,15 @@ mod _contextvars {
}

#[pymethod]
fn copy(&self) -> Self {
fn copy(&self, vm: &VirtualMachine) -> Self {
// Deep copy the vars - clone the underlying Hamt data, not just the PyRef
let vars_copy = HamtObject {
hamt: RefCell::new(self.inner.vars.hamt.borrow().clone()),
};
Self {
inner: ContextInner {
idx: Cell::new(usize::MAX),
vars: self.inner.vars.clone(),
vars: vars_copy.into_ref(&vm.ctx),
entered: Cell::new(false),
},
}
Expand Down Expand Up @@ -630,7 +634,7 @@ mod _contextvars {

#[pyfunction]
fn copy_context(vm: &VirtualMachine) -> PyContext {
PyContext::current(vm).copy()
PyContext::current(vm).copy(vm)
}

// Set Token.MISSING attribute
Expand Down
51 changes: 39 additions & 12 deletions crates/stdlib/src/socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,10 @@ mod _socket {
},
common::os::ErrorExt,
convert::{IntoPyException, ToPyObject, TryFromBorrowedObject, TryFromObject},
function::{ArgBytesLike, ArgMemoryBuffer, Either, FsPath, OptionalArg, OptionalOption},
function::{
ArgBytesLike, ArgMemoryBuffer, ArgStrOrBytesLike, Either, FsPath, OptionalArg,
OptionalOption,
},
types::{Constructor, DefaultConstructor, Initializer, Representable},
utils::ToCString,
};
Expand Down Expand Up @@ -2783,9 +2786,9 @@ mod _socket {
#[derive(FromArgs)]
struct GAIOptions {
#[pyarg(positional)]
host: Option<PyStrRef>,
host: Option<ArgStrOrBytesLike>,
#[pyarg(positional)]
port: Option<Either<PyStrRef, i32>>,
port: Option<Either<ArgStrOrBytesLike, i32>>,

#[pyarg(positional, default = c::AF_UNSPEC)]
family: i32,
Expand All @@ -2809,9 +2812,9 @@ mod _socket {
flags: opts.flags,
};

// Encode host using IDNA encoding
// Encode host: str uses IDNA encoding, bytes used as-is
let host_encoded: Option<String> = match opts.host.as_ref() {
Some(s) => {
Some(ArgStrOrBytesLike::Str(s)) => {
let encoded =
vm.state
.codec_registry
Expand All @@ -2820,19 +2823,43 @@ mod _socket {
.map_err(|_| vm.new_runtime_error("idna output is not utf8".to_owned()))?;
Some(host_str.to_owned())
}
Some(ArgStrOrBytesLike::Buf(b)) => {
let bytes = b.borrow_buf();
let host_str = core::str::from_utf8(&bytes).map_err(|_| {
vm.new_unicode_decode_error("host bytes is not utf8".to_owned())
})?;
Some(host_str.to_owned())
}
None => None,
};
let host = host_encoded.as_deref();

// Encode port using UTF-8
let port: Option<alloc::borrow::Cow<'_, str>> = match opts.port.as_ref() {
Some(Either::A(s)) => Some(alloc::borrow::Cow::Borrowed(s.to_str().ok_or_else(
|| vm.new_unicode_encode_error("surrogates not allowed".to_owned()),
)?)),
Some(Either::B(i)) => Some(alloc::borrow::Cow::Owned(i.to_string())),
// Encode port: str/bytes as service name, int as port number
let port_encoded: Option<String> = match opts.port.as_ref() {
Some(Either::A(sb)) => {
let port_str = match sb {
ArgStrOrBytesLike::Str(s) => {
// For str, check for surrogates and raise UnicodeEncodeError if found
s.to_str()
.ok_or_else(|| vm.new_unicode_encode_error("surrogates not allowed"))?
.to_owned()
}
ArgStrOrBytesLike::Buf(b) => {
// For bytes, check if it's valid UTF-8
let bytes = b.borrow_buf();
core::str::from_utf8(&bytes)
.map_err(|_| {
vm.new_unicode_decode_error("port is not utf8".to_owned())
})?
.to_owned()
}
};
Some(port_str)
}
Some(Either::B(i)) => Some(i.to_string()),
None => None,
};
let port = port.as_ref().map(|p| p.as_ref());
let port = port_encoded.as_deref();

let addrs = dns_lookup::getaddrinfo(host, port, Some(hints))
.map_err(|err| convert_socket_error(vm, err, SocketError::GaiError))?;
Expand Down
19 changes: 16 additions & 3 deletions crates/stdlib/src/ssl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ mod _ssl {
// Import error types used in this module (others are exposed via pymodule(with(...)))
use super::error::{
PySSLError, create_ssl_eof_error, create_ssl_want_read_error, create_ssl_want_write_error,
create_ssl_zero_return_error,
};
use alloc::sync::Arc;
use core::{
Expand Down Expand Up @@ -3593,7 +3594,7 @@ mod _ssl {
let mut conn_guard = self.connection.lock();
let conn = match conn_guard.as_mut() {
Some(conn) => conn,
None => return return_data(vec![], &buffer, vm),
None => return Err(create_ssl_zero_return_error(vm).upcast()),
};
use std::io::BufRead;
let mut reader = conn.reader();
Expand All @@ -3613,8 +3614,20 @@ mod _ssl {
return return_data(buf, &buffer, vm);
}
}
// Clean closure with close_notify - return empty data
return_data(vec![], &buffer, vm)
// Clean closure with close_notify
// CPython behavior depends on whether we've sent our close_notify:
// - If we've already sent close_notify (unwrap was called): raise SSLZeroReturnError
// - If we haven't sent close_notify yet: return empty bytes
let our_shutdown_state = *self.shutdown_state.lock();
if our_shutdown_state == ShutdownState::SentCloseNotify
|| our_shutdown_state == ShutdownState::Completed
{
// We already sent close_notify, now receiving peer's → SSLZeroReturnError
Err(create_ssl_zero_return_error(vm).upcast())
} else {
// We haven't sent close_notify yet → return empty bytes
return_data(vec![], &buffer, vm)
}
}
Err(crate::ssl::compat::SslError::WantRead) => {
// Non-blocking mode: would block
Expand Down
55 changes: 47 additions & 8 deletions crates/stdlib/src/ssl/compat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1552,6 +1552,11 @@ pub(super) fn ssl_read(

// Try to read plaintext from rustls buffer
if let Some(n) = try_read_plaintext(conn, buf)? {
if n == 0 {
// EOF from TLS - close_notify received
// Return ZeroReturn so Python raises SSLZeroReturnError
return Err(SslError::ZeroReturn);
}
return Ok(n);
}

Expand Down Expand Up @@ -1740,17 +1745,40 @@ pub(super) fn ssl_write(
let already_buffered = *socket.write_buffered_len.lock();

// Only write plaintext if not already buffered
// Track how much we wrote for partial write handling
let mut bytes_written_to_rustls = 0usize;

if already_buffered == 0 {
// Write plaintext to rustls (= SSL_write_ex internal buffer write)
{
bytes_written_to_rustls = {
let mut writer = conn.writer();
use std::io::Write;
writer
.write_all(data)
.map_err(|e| SslError::Syscall(format!("Write failed: {e}")))?;
}
// Mark data as buffered
*socket.write_buffered_len.lock() = data.len();
// Use write() instead of write_all() to support partial writes.
// In BIO mode (asyncio), when the internal buffer is full,
// we want to write as much as possible and return that count,
// rather than failing completely.
match writer.write(data) {
Ok(0) if !data.is_empty() => {
// Buffer is full and nothing could be written.
// In BIO mode, return WantWrite so the caller can
// drain the outgoing BIO and retry.
if is_bio {
return Err(SslError::WantWrite);
}
return Err(SslError::Syscall("Write failed: buffer full".to_string()));
}
Ok(n) => n,
Err(e) => {
if is_bio {
// In BIO mode, treat write errors as WantWrite
return Err(SslError::WantWrite);
}
return Err(SslError::Syscall(format!("Write failed: {e}")));
}
}
};
// Mark data as buffered (only the portion we actually wrote)
*socket.write_buffered_len.lock() = bytes_written_to_rustls;
} else if already_buffered != data.len() {
// Caller is retrying with different data - this is a protocol error
// Clear the buffer state and return an SSL error (bad write retry)
Expand Down Expand Up @@ -1826,10 +1854,21 @@ pub(super) fn ssl_write(
.map_err(SslError::Py)?;
}

// Determine how many bytes we actually wrote
let actual_written = if bytes_written_to_rustls > 0 {
// Fresh write: return what we wrote to rustls
bytes_written_to_rustls
} else if already_buffered > 0 {
// Retry of previous write: return the full buffered amount
already_buffered
} else {
data.len()
};

// Write completed successfully - clear buffer state
*socket.write_buffered_len.lock() = 0;

Ok(data.len())
Ok(actual_written)
}

// Helper functions (private-ish, used by public SSL functions)
Expand Down
8 changes: 6 additions & 2 deletions crates/vm/src/builtins/asyncgenerator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,12 @@ impl PyAsyncGen {
self.inner.frame().yield_from_target()
}
#[pygetset]
fn ag_frame(&self, _vm: &VirtualMachine) -> FrameRef {
self.inner.frame()
fn ag_frame(&self, _vm: &VirtualMachine) -> Option<FrameRef> {
if self.inner.closed() {
None
} else {
Some(self.inner.frame())
}
}
#[pygetset]
fn ag_running(&self, _vm: &VirtualMachine) -> bool {
Expand Down
8 changes: 6 additions & 2 deletions crates/vm/src/builtins/coroutine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,12 @@ impl PyCoroutine {
self.inner.frame().yield_from_target()
}
#[pygetset]
fn cr_frame(&self, _vm: &VirtualMachine) -> FrameRef {
self.inner.frame()
fn cr_frame(&self, _vm: &VirtualMachine) -> Option<FrameRef> {
if self.inner.closed() {
None
} else {
Some(self.inner.frame())
}
}
#[pygetset]
fn cr_running(&self, _vm: &VirtualMachine) -> bool {
Expand Down
Loading
Loading