make ord() work for any bytes-like

Author: timfel

graalpython/com.oracle.graal.python.test/src/tests/test_bytes.py

@@ -462,6 +462,18 @@ def test_endswith():
     assert not b.endswith(b"no")
 
 
+def test_ord():
+    b = b"123"
+    assert ord(b[0:1]) == 49
+    assert ord(bytearray(b)[0:1]) == 49
+    try:
+        ord(b)
+    except TypeError as e:
+        assert "expected a character" in str(e), str(e)
+    else:
+        assert False
+
+
 def test_find():
     b = b'mississippi'
     i = 105
@@ -538,7 +550,7 @@ def test_strip_bytes():
     assert b'abc'.rstrip(b'ac') == b'ab'
 
 class BaseTestSplit:
-    
+
     def test_string_error(self):
         self.assertRaises(TypeError, self.type2test(b'a b').split, ' ')
         self.assertRaises(TypeError, self.type2test(b'a b').rsplit, ' ')
@@ -601,10 +613,10 @@ def __init__(self, value):
                 self.value = value
             def __index__(self):
                 return self.value
-        
+
         self.assertEqual(self.type2test(b'ahoj jak\tse\nmas').split(maxsplit=MyIndexable(1)), [b'ahoj', b'jak\tse\nmas'])
         self.assertEqual(self.type2test(b'ahoj jak\tse\nmas').rsplit(maxsplit=MyIndexable(1)), [b'ahoj jak\tse', b'mas'])
-        
+
 class BytesSplitTest(BaseTestSplit, unittest.TestCase):
     type2test = bytes
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/BuiltinFunctions.java

@@ -91,6 +91,7 @@
 import com.oracle.graal.python.builtins.objects.PNotImplemented;
 import com.oracle.graal.python.builtins.objects.bytes.BytesNodes;
 import com.oracle.graal.python.builtins.objects.bytes.PBytes;
+import com.oracle.graal.python.builtins.objects.bytes.PIBytesLike;
 import com.oracle.graal.python.builtins.objects.code.PCode;
 import com.oracle.graal.python.builtins.objects.common.HashingCollectionNodes;
 import com.oracle.graal.python.builtins.objects.common.HashingStorageNodes;
@@ -1419,15 +1420,30 @@ public int ord(String chr) {
         }
 
         @Specialization
-        public int ord(PBytes chr,
+        public int ord(PIBytesLike chr,
                         @Cached("create()") SequenceStorageNodes.LenNode lenNode,
                         @Cached("create()") SequenceStorageNodes.GetItemNode getItemNode) {
             int len = lenNode.execute(chr.getSequenceStorage());
             if (len != 1) {
                 throw raise(TypeError, "ord() expected a character, but string of length %d found", len);
             }
 
-            return (byte) getItemNode.execute(chr.getSequenceStorage(), 0);
+            Object element = getItemNode.execute(chr.getSequenceStorage(), 0);
+            if (element instanceof Long) {
+                long e = (long) element;
+                if (e >= Byte.MIN_VALUE && e <= Byte.MAX_VALUE) {
+                    return (int) e;
+                }
+            } else if (element instanceof Integer) {
+                int e = (int) element;
+                if (e >= Byte.MIN_VALUE && e <= Byte.MAX_VALUE) {
+                    return e;
+                }
+            } else if (element instanceof Byte) {
+                return (byte) element;
+            }
+            CompilerDirectives.transferToInterpreter();
+            throw new IllegalStateException("got a bytes-like with non-byte elements");
         }
     }