diff --git a/package/PikaStdLib/PikaStdLib.pyi b/package/PikaStdLib/PikaStdLib.pyi index ff86aea40..da77a635b 100644 --- a/package/PikaStdLib/PikaStdLib.pyi +++ b/package/PikaStdLib/PikaStdLib.pyi @@ -28,7 +28,7 @@ class SysObj: def iter(arg: any) -> any: ... @staticmethod - def range(a1: int, a2: int) -> any: ... + def range(*ax) -> any: ... @staticmethod def print(*val, **ops): ... diff --git a/package/PikaStdLib/PikaStdLib_RangeObj.c b/package/PikaStdLib/PikaStdLib_RangeObj.c index a687144ca..2e4725b29 100644 --- a/package/PikaStdLib/PikaStdLib_RangeObj.c +++ b/package/PikaStdLib/PikaStdLib_RangeObj.c @@ -1,14 +1,10 @@ #include "PikaObj.h" Arg* PikaStdLib_RangeObj___next__(PikaObj* self) { - int a1 = args_getInt(self->list, "a1"); - int a2 = args_getInt(self->list, "a2"); - int start = 0; - int end = 0; - int foot = 1; + int start = args_getInt(self->list, "_start"); + int end = args_getInt(self->list, "_end"); + int step = args_getInt(self->list, "_step"); /* only start and end */ - start = a1; - end = a2; /* start */ int iter_i = args_getInt(self->list, "iter_i"); /* iter_i is not inited */ @@ -19,6 +15,6 @@ Arg* PikaStdLib_RangeObj___next__(PikaObj* self) { if (iter_i >= end) { return arg_newNull(); } - args_setInt(self->list, "iter_i", iter_i + foot); + args_setInt(self->list, "iter_i", iter_i + step); return arg_newInt(iter_i); } diff --git a/package/PikaStdLib/PikaStdLib_SysObj.c b/package/PikaStdLib/PikaStdLib_SysObj.c index 8c48facbd..2462f167d 100644 --- a/package/PikaStdLib/PikaStdLib_SysObj.c +++ b/package/PikaStdLib/PikaStdLib_SysObj.c @@ -184,11 +184,30 @@ Arg* PikaStdLib_SysObj_iter(PikaObj* self, Arg* arg) { return arg_newNull(); } -Arg* PikaStdLib_SysObj_range(PikaObj* self, int a1, int a2) { +Arg* PikaStdLib_SysObj_range(PikaObj* self, PikaTuple* ax) { /* set template arg to create rangeObj */ Arg* obj_arg = arg_newDirectObj(New_PikaStdLib_RangeObj); - obj_setInt(arg_getPtr(obj_arg), "a1", a1); - obj_setInt(arg_getPtr(obj_arg), "a2", a2); + PikaObj* range_obj = arg_getPtr(obj_arg); + if (tuple_getSize(ax) == 1) { + int start = 0; + int end = arg_getInt(tuple_getArg(ax, 0)); + obj_setInt(range_obj, "_start", start); + obj_setInt(range_obj, "_end", end); + obj_setInt(range_obj, "_step", 1); + } else if (tuple_getSize(ax) == 2) { + int start = arg_getInt(tuple_getArg(ax, 0)); + int end = arg_getInt(tuple_getArg(ax, 1)); + obj_setInt(range_obj, "_start", start); + obj_setInt(range_obj, "_end", end); + obj_setInt(range_obj, "_step", 1); + } else if (tuple_getSize(ax) == 3) { + int start = arg_getInt(tuple_getArg(ax, 0)); + int end = arg_getInt(tuple_getArg(ax, 1)); + int step = arg_getInt(tuple_getArg(ax, 2)); + obj_setInt(range_obj, "_start", start); + obj_setInt(range_obj, "_end", end); + obj_setInt(range_obj, "_step", step); + } return obj_arg; } diff --git a/port/linux/package/pikascript/PikaStdLib.pyi b/port/linux/package/pikascript/PikaStdLib.pyi index ff86aea40..da77a635b 100644 --- a/port/linux/package/pikascript/PikaStdLib.pyi +++ b/port/linux/package/pikascript/PikaStdLib.pyi @@ -28,7 +28,7 @@ class SysObj: def iter(arg: any) -> any: ... @staticmethod - def range(a1: int, a2: int) -> any: ... + def range(*ax) -> any: ... @staticmethod def print(*val, **ops): ... diff --git a/port/linux/package/pikascript/pikascript-lib/PikaStdLib/PikaStdLib_RangeObj.c b/port/linux/package/pikascript/pikascript-lib/PikaStdLib/PikaStdLib_RangeObj.c index a687144ca..2e4725b29 100644 --- a/port/linux/package/pikascript/pikascript-lib/PikaStdLib/PikaStdLib_RangeObj.c +++ b/port/linux/package/pikascript/pikascript-lib/PikaStdLib/PikaStdLib_RangeObj.c @@ -1,14 +1,10 @@ #include "PikaObj.h" Arg* PikaStdLib_RangeObj___next__(PikaObj* self) { - int a1 = args_getInt(self->list, "a1"); - int a2 = args_getInt(self->list, "a2"); - int start = 0; - int end = 0; - int foot = 1; + int start = args_getInt(self->list, "_start"); + int end = args_getInt(self->list, "_end"); + int step = args_getInt(self->list, "_step"); /* only start and end */ - start = a1; - end = a2; /* start */ int iter_i = args_getInt(self->list, "iter_i"); /* iter_i is not inited */ @@ -19,6 +15,6 @@ Arg* PikaStdLib_RangeObj___next__(PikaObj* self) { if (iter_i >= end) { return arg_newNull(); } - args_setInt(self->list, "iter_i", iter_i + foot); + args_setInt(self->list, "iter_i", iter_i + step); return arg_newInt(iter_i); } diff --git a/port/linux/package/pikascript/pikascript-lib/PikaStdLib/PikaStdLib_SysObj.c b/port/linux/package/pikascript/pikascript-lib/PikaStdLib/PikaStdLib_SysObj.c index 8c48facbd..2462f167d 100644 --- a/port/linux/package/pikascript/pikascript-lib/PikaStdLib/PikaStdLib_SysObj.c +++ b/port/linux/package/pikascript/pikascript-lib/PikaStdLib/PikaStdLib_SysObj.c @@ -184,11 +184,30 @@ Arg* PikaStdLib_SysObj_iter(PikaObj* self, Arg* arg) { return arg_newNull(); } -Arg* PikaStdLib_SysObj_range(PikaObj* self, int a1, int a2) { +Arg* PikaStdLib_SysObj_range(PikaObj* self, PikaTuple* ax) { /* set template arg to create rangeObj */ Arg* obj_arg = arg_newDirectObj(New_PikaStdLib_RangeObj); - obj_setInt(arg_getPtr(obj_arg), "a1", a1); - obj_setInt(arg_getPtr(obj_arg), "a2", a2); + PikaObj* range_obj = arg_getPtr(obj_arg); + if (tuple_getSize(ax) == 1) { + int start = 0; + int end = arg_getInt(tuple_getArg(ax, 0)); + obj_setInt(range_obj, "_start", start); + obj_setInt(range_obj, "_end", end); + obj_setInt(range_obj, "_step", 1); + } else if (tuple_getSize(ax) == 2) { + int start = arg_getInt(tuple_getArg(ax, 0)); + int end = arg_getInt(tuple_getArg(ax, 1)); + obj_setInt(range_obj, "_start", start); + obj_setInt(range_obj, "_end", end); + obj_setInt(range_obj, "_step", 1); + } else if (tuple_getSize(ax) == 3) { + int start = arg_getInt(tuple_getArg(ax, 0)); + int end = arg_getInt(tuple_getArg(ax, 1)); + int step = arg_getInt(tuple_getArg(ax, 2)); + obj_setInt(range_obj, "_start", start); + obj_setInt(range_obj, "_end", end); + obj_setInt(range_obj, "_step", step); + } return obj_arg; } diff --git a/port/linux/test/VM-test.cpp b/port/linux/test/VM-test.cpp index e953a6599..3be2083fa 100644 --- a/port/linux/test/VM-test.cpp +++ b/port/linux/test/VM-test.cpp @@ -1549,3 +1549,45 @@ TEST(vm, multi_return_fn) { EXPECT_EQ(pikaMemNow(), 0); } #endif + +TEST(vm, rang_1) { + /* init */ + pikaMemInfo.heapUsedMax = 0; + PikaObj* pikaMain = newRootObj("pikaMain", New_PikaMain); + extern unsigned char pikaModules_py_a[]; + obj_linkLibrary(pikaMain, pikaModules_py_a); + /* run */ + __platform_printf("BEGIN\r\n"); + obj_run(pikaMain, + "sum = 0\n" + "for i in range(10):\n" + " sum += i\n"); + /* collect */ + int sum = obj_getInt(pikaMain, "sum"); + /* assert */ + EXPECT_EQ(sum, 45); + /* deinit */ + obj_deinit(pikaMain); + EXPECT_EQ(pikaMemNow(), 0); +} + +TEST(vm, rang_3) { + /* init */ + pikaMemInfo.heapUsedMax = 0; + PikaObj* pikaMain = newRootObj("pikaMain", New_PikaMain); + extern unsigned char pikaModules_py_a[]; + obj_linkLibrary(pikaMain, pikaModules_py_a); + /* run */ + __platform_printf("BEGIN\r\n"); + obj_run(pikaMain, + "sum = 0\n" + "for i in range(-5, 20, 2):\n" + " sum += i\n"); + /* collect */ + int sum = obj_getInt(pikaMain, "sum"); + /* assert */ + EXPECT_EQ(sum, 91); + /* deinit */ + obj_deinit(pikaMain); + EXPECT_EQ(pikaMemNow(), 0); +}