diff --git a/crates/stdlib/src/csv.rs b/crates/stdlib/src/csv.rs index 792546a6d..c7b3d9a01 100644 --- a/crates/stdlib/src/csv.rs +++ b/crates/stdlib/src/csv.rs @@ -220,13 +220,8 @@ mod _csv { })?; let name: PyUtf8StrRef = name.try_into_utf8(vm)?; - let g = GLOBAL_HASHMAP.lock(); - - if let Some(dialect) = g.get(name.as_str()) { - return Ok(dialect.clone()); - } - - Err(new_csv_error(vm, "unknown dialect")) + let dialect = GLOBAL_HASHMAP.lock().get(name.as_str()).cloned(); + dialect.ok_or_else(|| new_csv_error(vm, "unknown dialect")) } #[pyfunction] @@ -239,9 +234,9 @@ mod _csv { })?; let name: PyUtf8StrRef = name.try_into_utf8(vm)?; - let mut g = GLOBAL_HASHMAP.lock(); + let removed = GLOBAL_HASHMAP.lock().remove(name.as_str()); - if let Some(_removed) = g.remove(name.as_str()) { + if removed.is_some() { return Ok(()); } @@ -250,10 +245,9 @@ mod _csv { #[pyfunction] fn list_dialects(vm: &VirtualMachine) -> rustpython_vm::builtins::PyListRef { - let g = GLOBAL_HASHMAP.lock(); - let t = g - .keys() - .cloned() + let names = GLOBAL_HASHMAP.lock().keys().cloned().collect_vec(); + let t = names + .into_iter() .map(|x| vm.ctx.new_str(x).into()) .collect_vec(); // .iter().map(|x| vm.ctx.new_str(x.clone()).into_pyobject(vm)).collect_vec(); @@ -268,14 +262,18 @@ mod _csv { #[pyfunction] fn field_size_limit(args: FieldSizeLimitArgs, vm: &VirtualMachine) -> PyResult { - let old_size = GLOBAL_FIELD_LIMIT.lock().to_owned(); if let OptionalArg::Present(limit) = args.new_limit { - let Ok(new_size) = limit.try_int(vm) else { - return Err(vm.new_type_error("limit must be an integer")); - }; - *GLOBAL_FIELD_LIMIT.lock() = new_size.try_to_primitive::(vm)?; + let new_size = limit + .downcast_ref_if_exact::(vm) + .ok_or_else(|| vm.new_type_error("limit must be an integer"))? + .try_to_primitive::(vm)?; + Ok(core::mem::replace( + &mut *GLOBAL_FIELD_LIMIT.lock(), + new_size, + )) + } else { + Ok(*GLOBAL_FIELD_LIMIT.lock()) } - Ok(old_size) } #[pyfunction] diff --git a/extra_tests/snippets/stdlib_csv.py b/extra_tests/snippets/stdlib_csv.py index 5911b0a1c..6bf29a67c 100644 --- a/extra_tests/snippets/stdlib_csv.py +++ b/extra_tests/snippets/stdlib_csv.py @@ -610,3 +610,32 @@ def test_dialect_lookup_and_conversion_order(): test_dialect_lookup_and_conversion_order() + + +def test_field_size_limit_types(): + index_calls = [] + + class IntSubclass(int): + pass + + class Index: + def __index__(self): + index_calls.append(True) + return 17 + + old_limit = csv.field_size_limit() + try: + assert csv.field_size_limit(23) == old_limit + for value in (True, IntSubclass(17), Index()): + with assert_raises(TypeError): + csv.field_size_limit(value) + assert csv.field_size_limit() == 23 + assert index_calls == [] + with assert_raises(OverflowError): + csv.field_size_limit(2**100) + assert csv.field_size_limit() == 23 + finally: + csv.field_size_limit(old_limit) + + +test_field_size_limit_types()