8000 feat: add array slicing (#18) (#20) · denosaurs/deno_python@610ca0f · GitHub
[go: up one dir, main page]
Skip to content

Commit 610ca0f

Browse files
tjosepotjosepo
authored
feat: add array slicing (#18) (#20)
1 parent c7d6bd4 commit 610ca0f

File tree

4 files changed

+187
-6
lines changed

4 files changed

+187
-6
lines changed

‎src/python.ts

Lines changed: 81 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// deno-lint-ignore-file no-explicit-any no-fallthrough
22
import { py } from "./ffi.ts";
3-
import { cstr } from "./util.ts";
3+
import { cstr, SliceItemRegExp } from "./util.ts";
44

55
/**
66
* Symbol used on proxied Python objects to point to the original PyObject object.
@@ -215,6 +215,17 @@ export class PyObject {
215215
}
216216
}
217217

218+
if (typeof name === "string" && isSlice(name)) {
219+
const slice = toSlice(name);
220+
const item = py.PyObject_GetItem(
221+
this.handle,
222+
slice.handle,
223+
) as Deno.UnsafePointer;
224+
if (item.value !== 0n) {
225+
return new PyObject(item).proxy;
226+
}
227+
}
228+
218229
// Don't wanna throw errors when accessing properties.
219230
const attr = this.maybeGetAttr(String(name))?.proxy;
220231

@@ -255,6 +266,14 @@ export class PyObject {
255266
PyObject.from(value).handle,
256267
);
257268
return true;
269+
} else if (isSlice(name)) {
270+
const slice = toSlice(name);
271+
py.PyObject_SetItem(
272+
this.handle,
273+
slice.handle,
274+
PyObject.from(value).handle,
275+
);
276+
return true;
258277
} else {
259278
return false;
260279
}
@@ -346,9 +365,8 @@ export class PyObject {
346365
} else {
347366
const dict = py.PyDict_New() as Deno.UnsafePointer;
348367
for (
349-
const [key, value] of (v instanceof Map
350-
? v.entries()
351-
: Object.entries(v))
368+
const [key, value]
369+
of (v instanceof Map ? v.entries() : Object.entries(v))
352370
) {
353371
const keyObj = PyObject.from(key);
354372
const valueObj = PyObject.from(value);
@@ -700,6 +718,8 @@ export class Python {
700718
tuple: any;
701719
/** Python `None` type proxied object */
702720
None: any;
721+
/** Python `Ellipsis` type proxied object */
722+
Ellipsis: any;
703723

704724
constructor() {
705725
py.Py_Initialize();
@@ -714,6 +734,7 @@ export class Python {
714734
this.bool = this.builtins.bool;
715735
this.set = this.builtins.set;
716736
this.tuple = this.builtins.tuple;
737+
this.Ellipsis = this.builtins.Ellipsis;
717738

718739
// Initialize arguments and executable path,
719740
// since some modules expect them to be set.
@@ -783,3 +804,59 @@ export class Python {
783804
* this object, such as `str`, `int`, `tuple`, etc.
784805
*/
785806
export const python = new Python();
807+
808+
/**
809+
* Returns true if the value can be converted into a Python slice or
810+
* slice tuple.
811+
*/
812+
function isSlice(value: unknown): boolean {
813+
if (typeof value !== "string") return false;
814+
if (!value.includes(":") && !value.includes("...")) return false;
815+
return value
816+
.split(",")
817+
.map((item) => (
818+
SliceItemRegExp.test(item) || // Slice
819+
/^\s*-?\d+\s*$/.test(item) || // Number
820+
/^\s*\.\.\.\s*$/.test(item) // Ellipsis
821+
))
822+
.reduce((a, b) => a && b, true);
823+
}
824+
825+
/**
826+
* Returns a PyObject that is either a slice or a tuple of slices.
827+
*/
828+
function toSlice(sliceList: string): PyObject {
829+
if (sliceList.includes(",")) {
830+
const pySlicesHandle = sliceList.split(",")
831+
.map(toSlice)
832+
.map((pyObject) => pyObject.handle);
833+
834+
const pyTuple_Pack = new Deno.UnsafeFnPointer(py.PyTuple_Pack, {
835+
parameters: ["i32", ...pySlicesHandle.map(() => "pointer" as const)],
836+
result: "pointer",
837+
});
838+
839+
const pyTupleHandle = pyTuple_Pack.call(
840+
pySlicesHandle.length,
841+
...pySlicesHandle,
842+
);
843+
return new PyObject(pyTupleHandle);
844+
} else if (/^\s*-?\d+\s*$/.test(sliceList)) {
845+
return PyObject.from(parseInt(sliceList));
846+
} else if (/^\s*\.\.\.\s*$/.test(sliceList)) {
847+
return PyObject.from(python.Ellipsis);
848+
} else {
849+
const [start, stop, step] = sliceList
850+
.split(":")
851+
.map((
852+
bound,
853+
) => (/^\s*-?\d+\s*$/.test(bound) ? parseInt(bound) : undefined));
854+
855+
const pySliceHandle = py.PySlice_New(
856+
PyObject.from(start).handle,
857+
PyObject.from(stop).handle,
858+
PyObject.from(step).handle,
859+
);
860+
return new PyObject(pySliceHandle);
861+
}
862+
}

‎src/symbols.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,4 +368,8 @@ export const SYMBOLS = {
368368
parameters: ["pointer", "i32"],
369369
result: "pointer",
370370
},
371+
372+
PyTuple_Pack: {
373+
type: "pointer",
374+
},
371375
} as const;

‎src/util.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,3 +52,11 @@ export function cstr(str: string): Uint8Array {
5252
encoder.encodeInto(str, buf);
5353
return buf;
5454
}
55+
56+
/**
57+
* Regular Expression used to test if a string is a `proper_slice`.
58+
*
59+
* Based on https://docs.python.org/3/reference/expressions.html#slicings
60+
*/
61+
export const SliceItemRegExp =
62+
/^\s*(-?\d+)?\s*:\s*(-?\d+)?\s*(:\s*(-?\d+)?\s*)?$/;

‎test/test.ts

Lines changed: 94 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,13 @@ Deno.test("types", async (t) => {
4343

4444
await t.step("dict", () => {
4545
const value = python.dict({ a: 1, b: 2 });
46-
assertEquals(value.valueOf(), new Map([["a", 1], ["b", 2]]));
46+
assertEquals(
47+
value.valueOf(),
48+
new Map([
49+
["a", 1],
50+
["b", 2],
51+
]),
52+
);
4753
});
4854

4955
await t.step("set", () => {
@@ -126,7 +132,9 @@ class Person:
126132
Deno.test("named argument", async (t) => {
127133
await t.step("single named argument", () => {
128134
assertEquals(
129-
python.str("Hello, {name}!").format(new NamedArgument("name", "world"))
135+
python
136+
.str("Hello, {name}!")
137+
.format(new NamedArgument("name", "world"))
130138
.valueOf(),
131139
"Hello, world!",
132140
);
@@ -171,3 +179,87 @@ Deno.test("custom proxy", () => {
171179
// Then, we use the wrapped proxy as if it were an original PyObject
172180
assertEquals(np.add(arr, 2).tolist().valueOf(), [3, 4, 5]);
173181
});
182+
183+
Deno.test("slice", async (t) => {
184+
await t.step("get", () => {
185+
const list = python.list([1, 2, 3, 4, 5, 6, 7, 8, 9]);
186+
assertEquals(list["1:"].valueOf(), [2, 3, 4, 5, 6, 7, 8, 9]);
187+
assertEquals(list["1:2"].valueOf(), [2]);
188+
assertEquals(list[":2"].valueOf(), [1< 179B span class=pl-kos>, 2]);
189+
assertEquals(list[":2:"].valueOf(), [1, 2]);
190+
assertEquals(list["0:3:2"].valueOf(), [1, 3]);
191+
assertEquals(list["-2:"].valueOf(), [8, 9]);
192+
assertEquals(list["::2"].valueOf(), [1, 3, 5, 7, 9]);
193+
});
194+
195+
await t.step("set", () => {
196+
const np = python.import("numpy");
197+
let list = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]);
198+
list["1:"] = -5;
199+
assertEquals(list.tolist().valueOf(), [1, -5, -5, -5, -5, -5, -5, -5, -5]);
200+
201+
list = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]);
202+
list["1::3"] = -5;
203+
assertEquals(list.tolist().valueOf(), [1, -5, 3, 4, -5, 6, 7, -5, 9]);
204+
205+
list = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]);
206+
list["1:2:3"] = -5;
207+
assertEquals(list.tolist().valueOf(), [1, -5, 3, 4, 5, 6, 7, 8, 9]);
208+
});
209+
});
210+
211+
Deno.test("slice list", async (t) => {
212+
const np = python.import("numpy");
213+
214+
await t.step("get", () => {
215+
const array = np.array([
216+
[1, 2, 3],
217+
[4, 5, 6],
218+
[7, 8, 9],
219+
]);
220+
assertEquals(array["0, :"].tolist().valueOf(), [1, 2, 3]);
221+
assertEquals(array["1:, ::2"].tolist().valueOf(), [
222+
[4, 6],
223+
[7, 9],
224+
]);
225+
assertEquals(array["1:, 0"].tolist().valueOf(), [4, 7]);
226+
});
227+
228+
await t.step("set", () => {
229+
const array = np.arange(15).reshape(3, 5);
230+
array["1:, ::2"] = -99;
231+
assertEquals(array.tolist().valueOf(), [
232+
[0, 1, 2, 3, 4],
233+
[-99, 6, -99, 8, -99],
234+
[-99, 11, -99, 13, -99],
235+
]);
236+
});
237+
238+
await t.step("whitespaces", () => {
239+
const array = np.array([
240+
[1, 2, 3],
241+
[4, 5, 6],
242+
[7, 8, 9],
243+
]);
244+
assertEquals(array[" 1 : , : : 2 "].tolist().valueOf(), [
245+
[4, 6],
246+
[7, 9],
247+
]);
248+
});
249+
250+
await t.step("3d slicing", () => {
251+
const a3 = np.array([[[10, 11, 12], [13, 14, 15], [16, 17, 18]], [
252+
[20, 21, 22],
253+
[23, 24, 25],
254+
[26, 27, 28],
255+
], [[30, 31, 32], [33, 34, 35], [36, 37, 38]]]);
256+
257+
assertEquals(a3["0, :, 1"].tolist().valueOf(), [11, 14, 17]);
258+
});
259+
260+
await t.step("ellipsis", () => {
261+
const a4 = np.arange(16).reshape(2, 2, 2, 2);
262+
263+
assertEquals(a4["1, ..., 1"].tolist().valueOf(), [[9, 11], [13, 15]]);
264+
});
265+
});

0 commit comments

Comments
 (0)
0