Skip to content

Commit ef2c0a6

Browse files
youknowoneyouknowone
authored
More FromArgs message (#6806)
1 parent 97167ab commit ef2c0a6

File tree

2 files changed

+39
-47
lines changed

2 files changed

+39
-47
lines changed

‎crates/stdlib/src/faulthandler.rs‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,10 @@ pub(crate) use decl::make_module;
44
#[pymodule(name = "faulthandler")]
55
mod decl {
66
use crate::vm::{
7-
PyObjectRef, PyResult, VirtualMachine, builtins::PyFloat, frame::Frame,
8-
function::OptionalArg, py_io::Write,
7+
PyObjectRef, PyResult, VirtualMachine,
8+
frame::Frame,
9+
function::{ArgIntoFloat, OptionalArg},
10+
py_io::Write,
911
};
1012
use alloc::sync::Arc;
1113
use core::sync::atomic::{AtomicBool, AtomicI32, Ordering};
@@ -762,8 +764,8 @@ mod decl {
762764
#[derive(FromArgs)]
763765
#[allow(unused)]
764766
struct DumpTracebackLaterArgs {
765-
#[pyarg(positional)]
766-
timeout: PyObjectRef,
767+
#[pyarg(positional, error_msg = "timeout must be a number (int or float)")]
768+
timeout: ArgIntoFloat,
767769
#[pyarg(any, default = false)]
768770
repeat: bool,
769771
#[pyarg(any, default)]
@@ -774,18 +776,7 @@ mod decl {
774776

775777
#[pyfunction]
776778
fn dump_traceback_later(args: DumpTracebackLaterArgs, vm: &VirtualMachine) -> PyResult<()> {
777-
use num_traits::ToPrimitive;
778-
// Convert timeout to f64 (accepting int or float)
779-
let timeout: f64 = if let Some(float) = args.timeout.downcast_ref::<PyFloat>() {
780-
float.to_f64()
781-
} else if let Some(int) = args.timeout.try_index_opt(vm).transpose()? {
782-
int.as_bigint()
783-
.to_i64()
784-
.ok_or_else(|| vm.new_overflow_error("timeout value is too large".to_owned()))?
785-
as f64
786-
} else {
787-
return Err(vm.new_type_error("timeout must be a number (int or float)".to_owned()));
788-
};
779+
let timeout: f64 = args.timeout.into_float();
789780

790781
if timeout <= 0.0 {
791782
return Err(vm.new_value_error("timeout must be greater than 0".to_owned()));

‎crates/stdlib/src/ssl.rs‎

Lines changed: 32 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,9 @@ mod _ssl {
4242
VirtualMachine,
4343
builtins::{PyBaseExceptionRef, PyBytesRef, PyListRef, PyStrRef, PyType, PyTypeRef},
4444
convert::IntoPyException,
45-
function::{ArgBytesLike, ArgMemoryBuffer, FuncArgs, OptionalArg, PyComparisonValue},
45+
function::{
46+
ArgBytesLike, ArgMemoryBuffer, Either, FuncArgs, OptionalArg, PyComparisonValue,
47+
},
4648
stdlib::warnings,
4749
types::{Comparable, Constructor, Hashable, PyComparisonOp, Representable},
4850
},
@@ -821,20 +823,20 @@ mod _ssl {
821823

822824
#[derive(FromArgs)]
823825
struct LoadVerifyLocationsArgs {
824-
#[pyarg(any, optional)]
825-
cafile: OptionalArg<Option<PyObjectRef>>,
826-
#[pyarg(any, optional)]
827-
capath: OptionalArg<Option<PyObjectRef>>,
828-
#[pyarg(any, optional)]
829-
cadata: OptionalArg<PyObjectRef>,
826+
#[pyarg(any, optional, error_msg = "path should be a str or bytes")]
827+
cafile: OptionalArg<Option<Either<PyStrRef, ArgBytesLike>>>,
828+
#[pyarg(any, optional, error_msg = "path should be a str or bytes")]
829+
capath: OptionalArg<Option<Either<PyStrRef, ArgBytesLike>>>,
830+
#[pyarg(any, optional, error_msg = "cadata should be a str or bytes")]
831+
cadata: OptionalArg<Option<Either<PyStrRef, ArgBytesLike>>>,
830832
}
831833

832834
#[derive(FromArgs)]
833835
struct LoadCertChainArgs {
834-
#[pyarg(any)]
835-
certfile: PyObjectRef,
836-
#[pyarg(any, optional)]
837-
keyfile: OptionalArg<Option<PyObjectRef>>,
836+
#[pyarg(any, error_msg = "path should be a str or bytes")]
837+
certfile: Either<PyStrRef, ArgBytesLike>,
838+
#[pyarg(any, optional, error_msg = "path should be a str or bytes")]
839+
keyfile: OptionalArg<Option<Either<PyStrRef, ArgBytesLike>>>,
838840
#[pyarg(any, optional)]
839841
password: OptionalArg<PyObjectRef>,
840842
}
@@ -1229,7 +1231,7 @@ mod _ssl {
12291231
// Check that at least one argument is provided
12301232
let has_cafile = matches!(&args.cafile, OptionalArg::Present(Some(_)));
12311233
let has_capath = matches!(&args.capath, OptionalArg::Present(Some(_)));
1232-
let has_cadata = matches!(&args.cadata, OptionalArg::Present(obj) if !vm.is_none(obj));
1234+
let has_cadata = matches!(&args.cadata, OptionalArg::Present(Some(_)));
12331235

12341236
if !has_cafile && !has_capath && !has_cadata {
12351237
return Err(
@@ -1250,10 +1252,8 @@ mod _ssl {
12501252
None
12511253
};
12521254

1253-
let cadata_parsed = if let OptionalArg::Present(ref cadata_obj) = args.cadata
1254-
&& !vm.is_none(cadata_obj)
1255-
{
1256-
let is_string = PyStrRef::try_from_object(vm, cadata_obj.clone()).is_ok();
1255+
let cadata_parsed = if let OptionalArg::Present(Some(ref cadata_obj)) = args.cadata {
1256+
let is_string = matches!(cadata_obj, Either::A(_));
12571257
let data_vec = self.parse_cadata_arg(cadata_obj, vm)?;
12581258
Some((data_vec, is_string))
12591259
} else {
@@ -1989,14 +1989,14 @@ mod _ssl {
19891989
// Helper functions (private):
19901990

19911991
/// Parse path argument (str or bytes) to string
1992-
fn parse_path_arg(arg: &PyObject, vm: &VirtualMachine) -> PyResult<String> {
1993-
if let Ok(s) = PyStrRef::try_from_object(vm, arg.to_owned()) {
1994-
Ok(s.as_str().to_owned())
1995-
} else if let Ok(b) = ArgBytesLike::try_from_object(vm, arg.to_owned()) {
1996-
String::from_utf8(b.borrow_buf().to_vec())
1997-
.map_err(|_| vm.new_value_error("path contains invalid UTF-8".to_owned()))
1998-
} else {
1999-
Err(vm.new_type_error("path should be a str or bytes".to_owned()))
1992+
fn parse_path_arg(
1993+
arg: &Either<PyStrRef, ArgBytesLike>,
1994+
vm: &VirtualMachine,
1995+
) -> PyResult<String> {
1996+
match arg {
1997+
Either::A(s) => Ok(s.as_str().to_owned()),
1998+
Either::B(b) => String::from_utf8(b.borrow_buf().to_vec())
1999+
.map_err(|_| vm.new_value_error("path contains invalid UTF-8".to_owned())),
20002000
}
20012001
}
20022002

@@ -2167,13 +2167,14 @@ mod _ssl {
21672167
}
21682168

21692169
/// Helper: Parse cadata argument (str or bytes)
2170-
fn parse_cadata_arg(&self, arg: &PyObject, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
2171-
if let Ok(s) = PyStrRef::try_from_object(vm, arg.to_owned()) {
2172-
Ok(s.as_str().as_bytes().to_vec())
2173-
} else if let Ok(b) = ArgBytesLike::try_from_object(vm, arg.to_owned()) {
2174-
Ok(b.borrow_buf().to_vec())
2175-
} else {
2176-
Err(vm.new_type_error("cadata should be a str or bytes".to_owned()))
2170+
fn parse_cadata_arg(
2171+
&self,
2172+
arg: &Either<PyStrRef, ArgBytesLike>,
2173+
_vm: &VirtualMachine,
2174+
) -> PyResult<Vec<u8>> {
2175+
match arg {
2176+
Either::A(s) => Ok(s.as_str().as_bytes().to_vec()),
2177+
Either::B(b) => Ok(b.borrow_buf().to_vec()),
21772178
}
21782179
}
21792180

0 commit comments

Comments
 (0)