Quellcode durchsuchen

support bytes slice

blueloveTH vor 1 Jahr
Ursprung
Commit
2219809fb4
2 geänderte Dateien mit 38 neuen und 5 gelöschten Zeilen
  1. 15 4
      src/pocketpy.cpp
  2. 23 1
      tests/11_bytes.py

+ 15 - 4
src/pocketpy.cpp

@@ -572,7 +572,7 @@ void init_builtins(VM* _vm) {
 #undef BIND_CMP_STR
 
     _vm->bind__getitem__(VM::tp_str, [](VM* vm, PyObject* _0, PyObject* _1) {
-        const Str& self = _CAST(Str&, _0);
+        const Str& self = PK_OBJ_GET(Str, _0);
         if(is_non_tagged_type(_1, vm->tp_slice)){
             const Slice& s = _CAST(Slice&, _1);
             int start, stop, step;
@@ -1090,9 +1090,20 @@ void init_builtins(VM* _vm) {
         return VAR(Bytes(buffer, list.size()));
     });
 
-    _vm->bind__getitem__(VM::tp_bytes, [](VM* vm, PyObject* obj, PyObject* index) {
-        const Bytes& self = _CAST(Bytes&, obj);
-        i64 i = CAST(i64, index);
+    _vm->bind__getitem__(VM::tp_bytes, [](VM* vm, PyObject* _0, PyObject* _1) {
+        const Bytes& self = PK_OBJ_GET(Bytes, _0);
+        if(is_non_tagged_type(_1, vm->tp_slice)){
+            const Slice& s = _CAST(Slice&, _1);
+            int start, stop, step;
+            vm->parse_int_slice(s, self.size(), start, stop, step);
+            int guess_max_size = abs(stop - start) / abs(step) + 1;
+            if(guess_max_size > self.size()) guess_max_size = self.size();
+            unsigned char* buffer = new unsigned char[guess_max_size];
+            int j = 0;      // actual size
+            PK_SLICE_LOOP(i, start, stop, step) buffer[j++] = self[i];
+            return VAR(Bytes(buffer, j));
+        }
+        i64 i = CAST(i64, _1);
         i = vm->normalized_index(i, self.size());
         return VAR(self[i]);
     });

+ 23 - 1
tests/11_bytes.py

@@ -11,4 +11,26 @@ assert b'\xff\xee' != b'1234'
 assert b'\xff\xee' == b'\xff\xee'
 
 a = '测试123'
-assert a == a.encode().decode()
+assert a == a.encode().decode()
+
+
+
+# test slice
+s = b"football"
+q = b"abcd"
+r = b"zoo"
+t = b"this is string example....wow!!!"
+assert s[0] == ord('f')
+assert s[1:4] == b'oot'
+assert s[:-1] == b'footbal'
+assert s[:10] == b'football'
+assert s[-3] == ord('a')
+assert t[-5:] == b'ow!!!'
+assert t[3:-3] == b's is string example....wow'
+
+a = b"Hello, World!"
+assert a[::-1] == b"!dlroW ,olleH"
+assert a[::2] == b"Hlo ol!"
+assert a[2:5:2] == b"lo"
+assert a[5:2:-1] == b",ol"
+assert a[5:2:-2] == b",l"