Skip to content

Commit 3085ead

Browse files
authored
Merge pull request RustPython#1055 from romab1998/feature/ord_bytes_bytearray
add bytearray and bytes input types for ord()
2 parents ba88716 + 7effca3 commit 3085ead

File tree

2 files changed

+38
-14
lines changed

2 files changed

+38
-14
lines changed

tests/snippets/builtin_ord.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,11 @@
33
assert ord("a") == 97
44
assert ord("é") == 233
55
assert ord("🤡") == 129313
6+
assert ord(b'a') == 97
7+
assert ord(bytearray(b'a')) == 97
68

79
assert_raises(TypeError, lambda: ord(), "ord() is called with no argument")
810
assert_raises(TypeError, lambda: ord(""), "ord() is called with an empty string")
911
assert_raises(TypeError, lambda: ord("ab"), "ord() is called with more than one character")
12+
assert_raises(TypeError, lambda: ord(b"ab"), "ord() expected a character, but string of length 2 found")
13+
assert_raises(TypeError, lambda: ord(1), "ord() expected a string, bytes or bytearray, but found int")

vm/src/builtins.rs

Lines changed: 34 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ use crate::pyobject::{
2828
};
2929
use crate::vm::VirtualMachine;
3030

31+
use crate::obj::objbyteinner::PyByteInner;
3132
#[cfg(not(target_arch = "wasm32"))]
3233
use crate::stdlib::io::io_open;
3334

@@ -530,20 +531,39 @@ fn builtin_oct(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
530531
}
531532

532533
fn builtin_ord(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
533-
arg_check!(vm, args, required = [(string, Some(vm.ctx.str_type()))]);
534-
let string = objstr::borrow_value(string);
535-
let string_len = string.chars().count();
536-
if string_len != 1 {
537-
return Err(vm.new_type_error(format!(
538-
"ord() expected a character, but string of length {} found",
539-
string_len
540-
)));
541-
}
542-
match string.chars().next() {
543-
Some(character) => Ok(vm.context().new_int(character as i32)),
544-
None => Err(vm.new_type_error(
545-
"ord() could not guess the integer representing this character".to_string(),
546-
)),
534+
arg_check!(vm, args, required = [(string, None)]);
535+
if objtype::isinstance(string, &vm.ctx.str_type()) {
536+
let string = objstr::borrow_value(string);
537+
let string_len = string.chars().count();
538+
if string_len != 1 {
539+
return Err(vm.new_type_error(format!(
540+
"ord() expected a character, but string of length {} found",
541+
string_len
542+
)));
543+
}
544+
match string.chars().next() {
545+
Some(character) => Ok(vm.context().new_int(character as i32)),
546+
None => Err(vm.new_type_error(
547+
"ord() could not guess the integer representing this character".to_string(),
548+
)),
549+
}
550+
} else if objtype::isinstance(string, &vm.ctx.bytearray_type())
551+
|| objtype::isinstance(string, &vm.ctx.bytes_type())
552+
{
553+
let inner = PyByteInner::try_from_object(vm, string.clone()).unwrap();
554+
let bytes_len = inner.elements.len();
555+
if bytes_len != 1 {
556+
return Err(vm.new_type_error(format!(
557+
"ord() expected a character, but string of length {} found",
558+
bytes_len
559+
)));
560+
}
561+
Ok(vm.context().new_int(inner.elements[0]))
562+
} else {
563+
Err(vm.new_type_error(format!(
564+
"ord() expected a string, bytes or bytearray, but found {}",
565+
string.class().name
566+
)))
547567
}
548568
}
549569

0 commit comments

Comments
 (0)