From 2219809fb435aaaff97fcc642e2dad2810754acf Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Sat, 6 Apr 2024 17:18:35 +0800
Subject: [PATCH] support bytes slice

---
 src/pocketpy.cpp  | 19 +++++++++++++++----
 tests/11_bytes.py | 24 +++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp
index a950c008..afc77631 100644
--- a/src/pocketpy.cpp
+++ b/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]);
     });
diff --git a/tests/11_bytes.py b/tests/11_bytes.py
index 76a247a6..bd854a6c 100644
--- a/tests/11_bytes.py
+++ b/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()
\ No newline at end of file
+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"
\ No newline at end of file