From 5a783d81f537acba54c1ff912ddd4f3e8795c5d9 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Sat, 14 Dec 2024 14:55:51 +0800
Subject: [PATCH] ...

---
 include/pocketpy/pocketpy.h |  8 ++++++++
 src/modules/pickle.c        | 10 +++++++++-
 src/public/py_dict.c        | 18 ++++++++++++++++++
 tests/90_pickle.py          |  1 +
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/include/pocketpy/pocketpy.h b/include/pocketpy/pocketpy.h
index 91fbff56..f4b40728 100644
--- a/include/pocketpy/pocketpy.h
+++ b/include/pocketpy/pocketpy.h
@@ -610,12 +610,20 @@ PK_API int py_dict_getitem(py_Ref self, py_Ref key) PY_RAISE PY_RETURN;
 PK_API bool py_dict_setitem(py_Ref self, py_Ref key, py_Ref val) PY_RAISE;
 /// -1: error, 0: not found, 1: found (and deleted)
 PK_API int py_dict_delitem(py_Ref self, py_Ref key) PY_RAISE;
+
 /// -1: error, 0: not found, 1: found
 PK_API int py_dict_getitem_by_str(py_Ref self, const char* key) PY_RAISE PY_RETURN;
+/// -1: error, 0: not found, 1: found
+PK_API int py_dict_getitem_by_int(py_Ref self, py_i64 key) PY_RAISE PY_RETURN;
 /// true: success, false: error
 PK_API bool py_dict_setitem_by_str(py_Ref self, const char* key, py_Ref val) PY_RAISE;
+/// true: success, false: error
+PK_API bool py_dict_setitem_by_int(py_Ref self, py_i64 key, py_Ref val) PY_RAISE;
 /// -1: error, 0: not found, 1: found (and deleted)
 PK_API int py_dict_delitem_by_str(py_Ref self, const char* key) PY_RAISE;
+/// -1: error, 0: not found, 1: found (and deleted)
+PK_API int py_dict_delitem_by_int(py_Ref self, py_i64 key) PY_RAISE;
+
 /// true: success, false: error
 PK_API bool
     py_dict_apply(py_Ref self, bool (*f)(py_Ref key, py_Ref val, void* ctx), void* ctx) PY_RAISE;
diff --git a/src/modules/pickle.c b/src/modules/pickle.c
index 6d820973..6a467335 100644
--- a/src/modules/pickle.c
+++ b/src/modules/pickle.c
@@ -8,7 +8,7 @@
 
 typedef enum {
     // clang-format off
-    PKL_NONE,
+    PKL_NONE, PKL_ELLIPSIS,
     PKL_INT8, PKL_INT16, PKL_INT32, PKL_INT64,
     PKL_FLOAT32, PKL_FLOAT64,
     PKL_TRUE, PKL_FALSE,
@@ -149,6 +149,10 @@ static bool pickle__write_object(PickleObject* buf, py_TValue* obj) {
             pkl__emit_op(buf, PKL_NONE);
             return true;
         }
+        case tp_ellipsis: {
+            pkl__emit_op(buf, PKL_ELLIPSIS);
+            return true;
+        }
         case tp_int: {
             py_i64 val = obj->_i64;
             pkl__emit_int(buf, val);
@@ -269,6 +273,10 @@ bool py_pickle_loads(const unsigned char* data, int size) {
                 py_pushnone();
                 break;
             }
+            case PKL_ELLIPSIS: {
+                py_newellipsis(py_pushtmp());
+                break;
+            }
             case PKL_INT8: {
                 int8_t val;
                 UNALIGNED_READ(&val, p);
diff --git a/src/public/py_dict.c b/src/public/py_dict.c
index 68ce3f71..c5523506 100644
--- a/src/public/py_dict.c
+++ b/src/public/py_dict.c
@@ -632,6 +632,24 @@ int py_dict_delitem_by_str(py_Ref self, const char* key) {
     return res;
 }
 
+int py_dict_getitem_by_int(py_Ref self, py_i64 key) {
+    py_TValue tmp;
+    py_newint(&tmp, key);
+    return py_dict_getitem(self, &tmp);
+}
+
+bool py_dict_setitem_by_int(py_Ref self, py_i64 key, py_Ref val) {
+    py_TValue tmp;
+    py_newint(&tmp, key);
+    return py_dict_setitem(self, &tmp, val);
+}
+
+int py_dict_delitem_by_int(py_Ref self, py_i64 key) {
+    py_TValue tmp;
+    py_newint(&tmp, key);
+    return py_dict_delitem(self, &tmp);
+}
+
 int py_dict_len(py_Ref self) {
     assert(py_isdict(self));
     Dict* ud = py_touserdata(self);
diff --git a/tests/90_pickle.py b/tests/90_pickle.py
index 34fb1bb4..65f92539 100644
--- a/tests/90_pickle.py
+++ b/tests/90_pickle.py
@@ -9,6 +9,7 @@ def test(data): # type: ignore
     assert data == o
 
 test(None)                      # PKL_NONE
+test(...)                       # PKL_ELLIPSIS
 test(1)                         # PKL_INT8
 test(277)                       # PKL_INT16
 test(-66666)                    # PKL_INT32