Zxc200611/pocketpy
1
0
Fork 0
mirror of https://github.com/pocketpy/pocketpy synced 2025-10-20 03:20:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

improve pickle

Browse Source
This commit is contained in:
blueloveTH 2024-12-15 20:57:14 +08:00
parent 3786abccdd
commit 79f18f8dc5
6 changed files with 286 additions and 139 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
include/pocketpy/config.h
View File

@ -36,6 +36,9 @@
// (not recommended to change this) // (not recommended to change this)
#define PK_MAX_CO_VARNAMES 64 #define PK_MAX_CO_VARNAMES 64
// This is the maximum character length of a module path
#define PK_MAX_MODULE_PATH_LEN 63
#ifdef _WIN32 #ifdef _WIN32
#define PK_PLATFORM_SEP '\\' #define PK_PLATFORM_SEP '\\'
#else #else

315
src/modules/pickle.c
View File

@ -24,44 +24,50 @@ typedef enum {
PKL_VEC2I, PKL_VEC3I, PKL_VEC2I, PKL_VEC3I,
PKL_TYPE, PKL_TYPE,
PKL_ARRAY2D, PKL_ARRAY2D,
PKL_TVALUE,
PKL_EOF, PKL_EOF,
// clang-format on // clang-format on
} PickleOp; } PickleOp;
typedef struct { typedef struct {
bool* used_types;
int used_types_length;
c11_smallmap_p2i memo; c11_smallmap_p2i memo;
c11_vector /*T=char*/ codes; c11_vector /*T=char*/ codes;
} PickleObject; } PickleObject;
typedef struct {
uint16_t memo_length;
} PickleObjectHeader;
static void PickleObject__ctor(PickleObject* self) { static void PickleObject__ctor(PickleObject* self) {
self->used_types_length = pk_current_vm->types.length;
self->used_types = malloc(self->used_types_length);
memset(self->used_types, 0, self->used_types_length);
c11_smallmap_p2i__ctor(&self->memo); c11_smallmap_p2i__ctor(&self->memo);
c11_vector__ctor(&self->codes, sizeof(char)); c11_vector__ctor(&self->codes, sizeof(char));
} }
static void PickleObject__dtor(PickleObject* self) { static void PickleObject__dtor(PickleObject* self) {
free(self->used_types);
c11_smallmap_p2i__dtor(&self->memo); c11_smallmap_p2i__dtor(&self->memo);
c11_vector__dtor(&self->codes); c11_vector__dtor(&self->codes);
} }
static bool PickleObject__py_submit(PickleObject* self, py_OutRef out) { static bool PickleObject__py_submit(PickleObject* self, py_OutRef out);
unsigned char* data = self->codes.data;
PickleObjectHeader* p =
(PickleObjectHeader*)py_newbytes(out, sizeof(PickleObjectHeader) + self->codes.length);
if(self->memo.length >= UINT16_MAX) c11__abort("PickleObject__py_submit(): memo overflow");
p->memo_length = (uint16_t)self->memo.length;
memcpy(p + 1, data, self->codes.length);
PickleObject__dtor(self);
return true;
}
static void PickleObject__write_bytes(PickleObject* buf, const void* data, int size) { static void PickleObject__write_bytes(PickleObject* buf, const void* data, int size) {
c11_vector__extend(char, &buf->codes, data, size); c11_vector__extend(char, &buf->codes, data, size);
} }
static void c11_sbuf__write_type_path(c11_sbuf* path_buf, py_Type type) {
py_TypeInfo* ti = pk__type_info(type);
if(py_isnil(&ti->module)) {
c11_sbuf__write_cstr(path_buf, py_name2str(ti->name));
return;
}
const char* mod_name = py_tostr(py_getdict(&ti->module, __name__));
c11_sbuf__write_cstr(path_buf, mod_name);
c11_sbuf__write_char(path_buf, '.');
c11_sbuf__write_cstr(path_buf, py_name2str(ti->name));
}
static void pkl__emit_op(PickleObject* buf, PickleOp op) { static void pkl__emit_op(PickleObject* buf, PickleOp op) {
c11_vector__push(char, &buf->codes, op); c11_vector__push(char, &buf->codes, op);
} }
@ -97,22 +103,10 @@ static py_i64 pkl__read_int(const unsigned char** p) {
(*p)++; (*p)++;
switch(op) { switch(op) {
// clang-format off // clang-format off
case PKL_INT_0: return 0; case PKL_INT_0: return 0; case PKL_INT_1: return 1; case PKL_INT_2: return 2; case PKL_INT_3: return 3;
case PKL_INT_1: return 1; case PKL_INT_4: return 4; case PKL_INT_5: return 5; case PKL_INT_6: return 6; case PKL_INT_7: return 7;
case PKL_INT_2: return 2; case PKL_INT_8: return 8; case PKL_INT_9: return 9; case PKL_INT_10: return 10; case PKL_INT_11: return 11;
case PKL_INT_3: return 3; case PKL_INT_12: return 12; case PKL_INT_13: return 13; case PKL_INT_14: return 14; case PKL_INT_15: return 15;
case PKL_INT_4: return 4;
case PKL_INT_5: return 5;
case PKL_INT_6: return 6;
case PKL_INT_7: return 7;
case PKL_INT_8: return 8;
case PKL_INT_9: return 9;
case PKL_INT_10: return 10;
case PKL_INT_11: return 11;
case PKL_INT_12: return 12;
case PKL_INT_13: return 13;
case PKL_INT_14: return 14;
case PKL_INT_15: return 15;
// clang-format on // clang-format on
case PKL_INT8: { case PKL_INT8: {
int8_t val; int8_t val;
@ -138,13 +132,6 @@ static py_i64 pkl__read_int(const unsigned char** p) {
} }
} }
const static char* pkl__read_cstr(const unsigned char** p) {
const char* p_str = (const char*)*p;
int length = strlen(p_str);
*p += length + 1; // include '\0'
return p_str;
}
static bool pickle_loads(int argc, py_Ref argv) { static bool pickle_loads(int argc, py_Ref argv) {
PY_CHECK_ARGC(1); PY_CHECK_ARGC(1);
PY_CHECK_ARG_TYPE(0, tp_bytes); PY_CHECK_ARG_TYPE(0, tp_bytes);
@ -165,11 +152,11 @@ void pk__add_module_pickle() {
py_bindfunc(mod, "dumps", pickle_dumps); py_bindfunc(mod, "dumps", pickle_dumps);
} }
static bool pickle__write_object(PickleObject* buf, py_TValue* obj); static bool pkl__write_object(PickleObject* buf, py_TValue* obj);
static bool pickle__write_array(PickleObject* buf, PickleOp op, py_TValue* arr, int length) { static bool pkl__write_array(PickleObject* buf, PickleOp op, py_TValue* arr, int length) {
for(int i = 0; i < length; i++) { for(int i = 0; i < length; i++) {
bool ok = pickle__write_object(buf, arr + i); bool ok = pkl__write_object(buf, arr + i);
if(!ok) return false; if(!ok) return false;
} }
pkl__emit_op(buf, op); pkl__emit_op(buf, op);
@ -177,36 +164,44 @@ static bool pickle__write_array(PickleObject* buf, PickleOp op, py_TValue* arr,
return true; return true;
} }
static bool pickle__write_dict_kv(py_Ref k, py_Ref v, void* ctx) { static bool pkl__write_dict_kv(py_Ref k, py_Ref v, void* ctx) {
PickleObject* buf = (PickleObject*)ctx; PickleObject* buf = (PickleObject*)ctx;
if(!pickle__write_object(buf, k)) return false; if(!pkl__write_object(buf, k)) return false;
if(!pickle__write_object(buf, v)) return false; if(!pkl__write_object(buf, v)) return false;
return true; return true;
} }
static bool pickle__write_object(PickleObject* buf, py_TValue* obj) { static bool pkl__try_memo(PickleObject* buf, PyObject* memo_key) {
if(obj->is_ptr) {
void* memo_key = obj->_obj;
int index = c11_smallmap_p2i__get(&buf->memo, memo_key, -1); int index = c11_smallmap_p2i__get(&buf->memo, memo_key, -1);
if(index != -1) { if(index != -1) {
pkl__emit_op(buf, PKL_MEMO_GET); pkl__emit_op(buf, PKL_MEMO_GET);
pkl__emit_int(buf, index); pkl__emit_int(buf, index);
return true; return true;
} }
} return false;
}
static void pkl__store_memo(PickleObject* buf, PyObject* memo_key) {
int index = buf->memo.length;
c11_smallmap_p2i__set(&buf->memo, memo_key, index);
pkl__emit_op(buf, PKL_MEMO_SET);
pkl__emit_int(buf, index);
}
static bool pkl__write_object(PickleObject* buf, py_TValue* obj) {
switch(obj->type) { switch(obj->type) {
case tp_NoneType: { case tp_NoneType: {
pkl__emit_op(buf, PKL_NONE); pkl__emit_op(buf, PKL_NONE);
break; return true;
} }
case tp_ellipsis: { case tp_ellipsis: {
pkl__emit_op(buf, PKL_ELLIPSIS); pkl__emit_op(buf, PKL_ELLIPSIS);
break; return true;
} }
case tp_int: { case tp_int: {
py_i64 val = obj->_i64; py_i64 val = obj->_i64;
pkl__emit_int(buf, val); pkl__emit_int(buf, val);
break; return true;
} }
case tp_float: { case tp_float: {
py_f64 val = obj->_f64; py_f64 val = obj->_f64;
@ -218,64 +213,90 @@ static bool pickle__write_object(PickleObject* buf, py_TValue* obj) {
pkl__emit_op(buf, PKL_FLOAT64); pkl__emit_op(buf, PKL_FLOAT64);
PickleObject__write_bytes(buf, &val, 8); PickleObject__write_bytes(buf, &val, 8);
} }
break; return true;
} }
case tp_bool: { case tp_bool: {
bool val = obj->_bool; bool val = obj->_bool;
pkl__emit_op(buf, val ? PKL_TRUE : PKL_FALSE); pkl__emit_op(buf, val ? PKL_TRUE : PKL_FALSE);
break; return true;
} }
case tp_str: { case tp_str: {
if(pkl__try_memo(buf, obj->_obj))
return true;
else {
pkl__emit_op(buf, PKL_STRING); pkl__emit_op(buf, PKL_STRING);
c11_sv sv = py_tosv(obj); c11_sv sv = py_tosv(obj);
pkl__emit_int(buf, sv.size); pkl__emit_int(buf, sv.size);
PickleObject__write_bytes(buf, sv.data, sv.size); PickleObject__write_bytes(buf, sv.data, sv.size);
break; }
pkl__store_memo(buf, obj->_obj);
return true;
} }
case tp_bytes: { case tp_bytes: {
if(pkl__try_memo(buf, obj->_obj))
return true;
else {
pkl__emit_op(buf, PKL_BYTES); pkl__emit_op(buf, PKL_BYTES);
int size; int size;
unsigned char* data = py_tobytes(obj, &size); unsigned char* data = py_tobytes(obj, &size);
pkl__emit_int(buf, size); pkl__emit_int(buf, size);
PickleObject__write_bytes(buf, data, size); PickleObject__write_bytes(buf, data, size);
break; }
pkl__store_memo(buf, obj->_obj);
return true;
} }
case tp_list: { case tp_list: {
bool ok = pickle__write_array(buf, PKL_BUILD_LIST, py_list_data(obj), py_list_len(obj)); if(pkl__try_memo(buf, obj->_obj))
return true;
else {
bool ok =
pkl__write_array(buf, PKL_BUILD_LIST, py_list_data(obj), py_list_len(obj));
if(!ok) return false; if(!ok) return false;
break; }
pkl__store_memo(buf, obj->_obj);
return true;
} }
case tp_tuple: { case tp_tuple: {
if(pkl__try_memo(buf, obj->_obj))
return true;
else {
bool ok = bool ok =
pickle__write_array(buf, PKL_BUILD_TUPLE, py_tuple_data(obj), py_tuple_len(obj)); pkl__write_array(buf, PKL_BUILD_TUPLE, py_tuple_data(obj), py_tuple_len(obj));
if(!ok) return false; if(!ok) return false;
break; }
pkl__store_memo(buf, obj->_obj);
return true;
} }
case tp_dict: { case tp_dict: {
bool ok = py_dict_apply(obj, pickle__write_dict_kv, (void*)buf); if(pkl__try_memo(buf, obj->_obj))
return true;
else {
bool ok = py_dict_apply(obj, pkl__write_dict_kv, (void*)buf);
if(!ok) return false; if(!ok) return false;
pkl__emit_op(buf, PKL_BUILD_DICT); pkl__emit_op(buf, PKL_BUILD_DICT);
pkl__emit_int(buf, py_dict_len(obj)); pkl__emit_int(buf, py_dict_len(obj));
break; }
pkl__store_memo(buf, obj->_obj);
return true;
} }
case tp_vec2: { case tp_vec2: {
c11_vec2 val = py_tovec2(obj); c11_vec2 val = py_tovec2(obj);
pkl__emit_op(buf, PKL_VEC2); pkl__emit_op(buf, PKL_VEC2);
PickleObject__write_bytes(buf, &val, sizeof(c11_vec2)); PickleObject__write_bytes(buf, &val, sizeof(c11_vec2));
break; return true;
} }
case tp_vec3: { case tp_vec3: {
c11_vec3 val = py_tovec3(obj); c11_vec3 val = py_tovec3(obj);
pkl__emit_op(buf, PKL_VEC3); pkl__emit_op(buf, PKL_VEC3);
PickleObject__write_bytes(buf, &val, sizeof(c11_vec3)); PickleObject__write_bytes(buf, &val, sizeof(c11_vec3));
break; return true;
} }
case tp_vec2i: { case tp_vec2i: {
c11_vec2i val = py_tovec2i(obj); c11_vec2i val = py_tovec2i(obj);
pkl__emit_op(buf, PKL_VEC2I); pkl__emit_op(buf, PKL_VEC2I);
pkl__emit_int(buf, val.x); pkl__emit_int(buf, val.x);
pkl__emit_int(buf, val.y); pkl__emit_int(buf, val.y);
break; return true;
} }
case tp_vec3i: { case tp_vec3i: {
c11_vec3i val = py_tovec3i(obj); c11_vec3i val = py_tovec3i(obj);
@ -283,53 +304,51 @@ static bool pickle__write_object(PickleObject* buf, py_TValue* obj) {
pkl__emit_int(buf, val.x); pkl__emit_int(buf, val.x);
pkl__emit_int(buf, val.y); pkl__emit_int(buf, val.y);
pkl__emit_int(buf, val.z); pkl__emit_int(buf, val.z);
break; return true;
} }
case tp_type: { case tp_type: {
pkl__emit_op(buf, PKL_TYPE); pkl__emit_op(buf, PKL_TYPE);
py_TypeInfo* ti = pk__type_info(py_totype(obj)); py_Type type = py_totype(obj);
const char* mod_name = py_tostr(py_getdict(&ti->module, __name__)); buf->used_types[type] = true;
c11_sbuf path_buf; pkl__emit_int(buf, type);
c11_sbuf__ctor(&path_buf); return true;
c11_sbuf__write_cstr(&path_buf, mod_name);
c11_sbuf__write_cstr(&path_buf, "@");
c11_sbuf__write_cstr(&path_buf, py_name2str(ti->name));
c11_string* path = c11_sbuf__submit(&path_buf);
// include '\0'
PickleObject__write_bytes(buf, path->data, path->size + 1);
c11_string__delete(path);
break;
} }
case tp_array2d: { case tp_array2d: {
if(pkl__try_memo(buf, obj->_obj))
return true;
else {
c11_array2d* arr = py_touserdata(obj); c11_array2d* arr = py_touserdata(obj);
for(int i = 0; i < arr->numel; i++) { for(int i = 0; i < arr->numel; i++) {
if(arr->data[i].is_ptr) if(arr->data[i].is_ptr)
return TypeError( return TypeError(
"'array2d' object is not picklable because it contains heap-allocated objects"); "'array2d' object is not picklable because it contains heap-allocated objects");
buf->used_types[arr->data[i].type] = true;
} }
pkl__emit_op(buf, PKL_ARRAY2D); pkl__emit_op(buf, PKL_ARRAY2D);
pkl__emit_int(buf, arr->n_cols); pkl__emit_int(buf, arr->n_cols);
pkl__emit_int(buf, arr->n_rows); pkl__emit_int(buf, arr->n_rows);
// TODO: fix type index which is not stable
PickleObject__write_bytes(buf, arr->data, arr->numel * sizeof(py_TValue)); PickleObject__write_bytes(buf, arr->data, arr->numel * sizeof(py_TValue));
break;
}
default: return TypeError("'%t' object is not picklable", obj->type);
}
if(obj->is_ptr) {
void* memo_key = obj->_obj;
int index = buf->memo.length;
c11_smallmap_p2i__set(&buf->memo, memo_key, index);
pkl__emit_op(buf, PKL_MEMO_SET);
pkl__emit_int(buf, index);
} }
pkl__store_memo(buf, obj->_obj);
return true; return true;
}
default: {
if(!obj->is_ptr) {
pkl__emit_op(buf, PKL_TVALUE);
PickleObject__write_bytes(buf, obj, sizeof(py_TValue));
buf->used_types[obj->type] = true;
return true;
}
return TypeError("'%t' object is not picklable", obj->type);
}
}
c11__unreachable();
} }
bool py_pickle_dumps(py_Ref val) { bool py_pickle_dumps(py_Ref val) {
PickleObject buf; PickleObject buf;
PickleObject__ctor(&buf); PickleObject__ctor(&buf);
bool ok = pickle__write_object(&buf, val); bool ok = pkl__write_object(&buf, val);
if(!ok) { if(!ok) {
PickleObject__dtor(&buf); PickleObject__dtor(&buf);
return false; return false;
@ -338,12 +357,74 @@ bool py_pickle_dumps(py_Ref val) {
return PickleObject__py_submit(&buf, py_retval()); return PickleObject__py_submit(&buf, py_retval());
} }
static py_Type pkl__header_find_type(c11_sv path) {
int sep_index = c11_sv__rindex(path, '.');
if(sep_index == -1) return py_gettype(NULL, py_namev(path));
c11_sv mod_name = c11_sv__slice2(path, 0, sep_index);
c11_sv name = c11_sv__slice(path, sep_index + 1);
char buf[PK_MAX_MODULE_PATH_LEN + 1];
memcpy(buf, mod_name.data, mod_name.size);
buf[mod_name.size] = '\0';
return py_gettype(buf, py_namev(name));
}
static c11_sv pkl__header_read_sv(const unsigned char** p, char sep) {
c11_sv text;
text.data = (const char*)*p;
const char* p_end = strchr(text.data, sep);
assert(p_end != NULL);
text.size = p_end - text.data;
*p = (const unsigned char*)p_end + 1;
return text;
}
static py_i64 pkl__header_read_int(const unsigned char** p, char sep) {
c11_sv text = pkl__header_read_sv(p, sep);
py_i64 out;
IntParsingResult res = c11__parse_uint(text, &out, 10);
assert(res == IntParsing_SUCCESS);
return out;
}
bool py_pickle_loads_body(const unsigned char* p, int memo_length, c11_smallmap_n2i* type_mapping);
bool py_pickle_loads(const unsigned char* data, int size) { bool py_pickle_loads(const unsigned char* data, int size) {
PickleObjectHeader* header = (PickleObjectHeader*)data; const unsigned char* p = data;
const unsigned char* p = (const unsigned char*)(header + 1);
c11_smallmap_n2i type_mapping;
c11_smallmap_n2i__ctor(&type_mapping);
while(true) {
if(*p == '\n') {
p++;
break;
}
py_Type type = pkl__header_read_int(&p, '(');
c11_sv path = pkl__header_read_sv(&p, ')');
py_Type new_type = pkl__header_find_type(path);
if(new_type == 0) {
c11_smallmap_n2i__dtor(&type_mapping);
return ImportError("cannot find type '%v'", path);
}
if(type != new_type) c11_smallmap_n2i__set(&type_mapping, type, new_type);
}
int memo_length = pkl__header_read_int(&p, '\n');
bool ok = py_pickle_loads_body(p, memo_length, &type_mapping);
c11_smallmap_n2i__dtor(&type_mapping);
return ok;
}
static py_Type pkl__fix_type(py_Type type, c11_smallmap_n2i* type_mapping) {
int new_type = c11_smallmap_n2i__get(type_mapping, type, -1);
if(new_type != -1) return (py_Type)new_type;
return type;
}
bool py_pickle_loads_body(const unsigned char* p, int memo_length, c11_smallmap_n2i* type_mapping) {
py_StackRef p0 = py_peek(0); py_StackRef p0 = py_peek(0);
py_StackRef memo = py_pushtmp(); py_StackRef memo = py_pushtmp();
py_newtuple(memo, header->memo_length); py_newtuple(memo, memo_length);
while(true) { while(true) {
PickleOp op = (PickleOp)*p; PickleOp op = (PickleOp)*p;
p++; p++;
@ -504,18 +585,9 @@ bool py_pickle_loads(const unsigned char* data, int size) {
break; break;
} }
case PKL_TYPE: { case PKL_TYPE: {
const char* path = pkl__read_cstr(&p); py_Type type = (py_Type)pkl__read_int(&p);
char* sep_index = strchr(path, '@'); type = pkl__fix_type(type, type_mapping);
assert(sep_index != NULL); py_push(py_tpobject(type));
*sep_index = '\0';
const char* mod_name = path;
const char* type_name = sep_index + 1;
py_Type t = py_gettype(mod_name, py_name(type_name));
*sep_index = '@';
if(t == 0) {
return ImportError("cannot import '%s' from '%s'", type_name, mod_name);
}
py_push(py_tpobject(t));
break; break;
} }
case PKL_ARRAY2D: { case PKL_ARRAY2D: {
@ -524,9 +596,19 @@ bool py_pickle_loads(const unsigned char* data, int size) {
c11_array2d* arr = py_newarray2d(py_pushtmp(), n_cols, n_rows); c11_array2d* arr = py_newarray2d(py_pushtmp(), n_cols, n_rows);
int total_size = arr->numel * sizeof(py_TValue); int total_size = arr->numel * sizeof(py_TValue);
memcpy(arr->data, p, total_size); memcpy(arr->data, p, total_size);
for(int i = 0; i < arr->numel; i++) {
arr->data[i].type = pkl__fix_type(arr->data[i].type, type_mapping);
}
p += total_size; p += total_size;
break; break;
} }
case PKL_TVALUE: {
py_TValue* tmp = py_pushtmp();
memcpy(tmp, p, sizeof(py_TValue));
tmp->type = pkl__fix_type(tmp->type, type_mapping);
p += sizeof(py_TValue);
break;
}
case PKL_EOF: { case PKL_EOF: {
// [memo, obj] // [memo, obj]
if(py_peek(0) - p0 != 2) return ValueError("invalid pickle data"); if(py_peek(0) - p0 != 2) return ValueError("invalid pickle data");
@ -537,6 +619,33 @@ bool py_pickle_loads(const unsigned char* data, int size) {
default: c11__unreachable(); default: c11__unreachable();
} }
} }
c11__unreachable();
}
static bool PickleObject__py_submit(PickleObject* self, py_OutRef out) {
c11_sbuf cleartext;
c11_sbuf__ctor(&cleartext);
// line 1: type mappping
for(py_Type type = 0; type < self->used_types_length; type++) {
if(self->used_types[type]) {
c11_sbuf__write_int(&cleartext, type);
c11_sbuf__write_char(&cleartext, '(');
c11_sbuf__write_type_path(&cleartext, type);
c11_sbuf__write_char(&cleartext, ')');
}
}
c11_sbuf__write_char(&cleartext, '\n');
// line 2: memo length
c11_sbuf__write_int(&cleartext, self->memo.length);
c11_sbuf__write_char(&cleartext, '\n');
// -------------------------------------------------- //
c11_string* header = c11_sbuf__submit(&cleartext);
int total_size = header->size + self->codes.length;
unsigned char* p = py_newbytes(py_retval(), total_size);
memcpy(p, header->data, header->size);
memcpy(p + header->size, self->codes.data, self->codes.length);
c11_string__delete(header);
PickleObject__dtor(self);
return true; return true;
} }

9
src/public/cast.c
View File

@ -24,7 +24,7 @@ bool py_castfloat(py_Ref self, double* out) {
} }
} }
bool py_castfloat32(py_Ref self, float *out){ bool py_castfloat32(py_Ref self, float* out) {
switch(self->type) { switch(self->type) {
case tp_int: *out = (float)self->_i64; return true; case tp_int: *out = (float)self->_i64; return true;
case tp_float: *out = (float)self->_f64; return true; case tp_float: *out = (float)self->_f64; return true;
@ -77,8 +77,13 @@ bool py_issubclass(py_Type derived, py_Type base) {
py_Type py_typeof(py_Ref self) { return self->type; } py_Type py_typeof(py_Ref self) { return self->type; }
py_Type py_gettype(const char* module, py_Name name) { py_Type py_gettype(const char* module, py_Name name) {
py_Ref mod = py_getmodule(module); py_Ref mod;
if(module != NULL) {
mod = py_getmodule(module);
if(!mod) return 0; if(!mod) return 0;
} else {
mod = &pk_current_vm->builtins;
}
py_Ref object = py_getdict(mod, name); py_Ref object = py_getdict(mod, name);
if(object && py_istype(object, tp_type)) return py_totype(object); if(object && py_istype(object, tp_type)) return py_totype(object);
return 0; return 0;

1
src/public/modules.c
View File

@ -23,6 +23,7 @@ void py_setglobal(py_Name name, py_Ref val) { py_setdict(&pk_current_vm->main, n
py_Ref py_newmodule(const char* path) { py_Ref py_newmodule(const char* path) {
ManagedHeap* heap = &pk_current_vm->heap; ManagedHeap* heap = &pk_current_vm->heap;
if(strlen(path) > PK_MAX_MODULE_PATH_LEN) c11__abort("module path too long: %s", path);
py_Ref r0 = py_pushtmp(); py_Ref r0 = py_pushtmp();
py_Ref r1 = py_pushtmp(); py_Ref r1 = py_pushtmp();

35
tests/90_pickle.py
View File

@ -33,9 +33,9 @@ test(vec3i) # PKL_TYPE
print('-'*50) print('-'*50)
from array2d import array2d from array2d import array2d
a = array2d[int].fromlist([ a = array2d[int | bool | vec2i].fromlist([
[1, 2, 3], [1, 2, vec2i.LEFT],
[4, 5, 6] [4, True, 6]
]) ])
a_encoded = pkl.dumps(a) a_encoded = pkl.dumps(a)
print(a_encoded) print(a_encoded)
@ -72,6 +72,35 @@ assert b is not a
assert b[0] is b[2] assert b[0] is b[2]
assert b[1] is b[3] assert b[1] is b[3]
from pkpy import TValue
class Base(TValue[int]):
def __eq__(self, other):
return self.value == other.value
def __ne__(self, other):
return self.value != other.value
class TVal(Base): pass # type: ignore
test(TVal(1))
old_bytes = pkl.dumps(TVal(1))
print(old_bytes)
# re-define the class so it will have a new type id
class TVal(Base): pass
# see if we can still load the old data
decoded = pkl.loads(old_bytes)
assert decoded == TVal(1)
print(pkl.dumps(decoded))
# test array2d with TValue
a = array2d[TVal].fromlist([
[TVal(1), TVal(2)],
[TVal(3), 1]])
test(a)
exit() exit()
from pickle import dumps, loads, _wrap, _unwrap from pickle import dumps, loads, _wrap, _unwrap

2
tests/98_lz4.py
View File

@ -26,6 +26,6 @@ for i in range(100):
ratio = test(gen_data()) ratio = test(gen_data())
# print(f'compression ratio: {ratio:.2f}') # print(f'compression ratio: {ratio:.2f}')
# test 1GB random data # test 64MB random data (require 1GB list[int] buffer)
rnd = [random.randint(0, 255) for _ in range(1024*1024*1024//16)] rnd = [random.randint(0, 255) for _ in range(1024*1024*1024//16)]
test(bytes(rnd)) test(bytes(rnd))