feat: add array slicing (#18) (#20)

denosaurs · Apr 27, 2022 · 610ca0f · 610ca0f
1 parent c7d6bd4
commit 610ca0f
Show file tree

Hide file tree

Showing 4 changed files with 187 additions and 6 deletions.
diff --git a/src/python.ts b/src/python.ts
@@ -1,6 +1,6 @@
 // deno-lint-ignore-file no-explicit-any no-fallthrough
 import { py } from "./ffi.ts";
-import { cstr } from "./util.ts";
+import { cstr, SliceItemRegExp } from "./util.ts";
 
 /**
  * Symbol used on proxied Python objects to point to the original PyObject object.
@@ -215,6 +215,17 @@ export class PyObject {
           }
         }
 
+        if (typeof name === "string" && isSlice(name)) {
+          const slice = toSlice(name);
+          const item = py.PyObject_GetItem(
+            this.handle,
+            slice.handle,
+          ) as Deno.UnsafePointer;
+          if (item.value !== 0n) {
+            return new PyObject(item).proxy;
+          }
+        }
+
         // Don't wanna throw errors when accessing properties.
         const attr = this.maybeGetAttr(String(name))?.proxy;
 
@@ -255,6 +266,14 @@ export class PyObject {
             PyObject.from(value).handle,
           );
           return true;
+        } else if (isSlice(name)) {
+          const slice = toSlice(name);
+          py.PyObject_SetItem(
+            this.handle,
+            slice.handle,
+            PyObject.from(value).handle,
+          );
+          return true;
         } else {
           return false;
         }
@@ -346,9 +365,8 @@ export class PyObject {
         } else {
           const dict = py.PyDict_New() as Deno.UnsafePointer;
           for (
-            const [key, value] of (v instanceof Map
-              ? v.entries()
-              : Object.entries(v))
+            const [key, value]
+              of (v instanceof Map ? v.entries() : Object.entries(v))
           ) {
             const keyObj = PyObject.from(key);
             const valueObj = PyObject.from(value);
@@ -700,6 +718,8 @@ export class Python {
   tuple: any;
   /** Python `None` type proxied object */
   None: any;
+  /** Python `Ellipsis` type proxied object */
+  Ellipsis: any;
 
   constructor() {
     py.Py_Initialize();
@@ -714,6 +734,7 @@ export class Python {
     this.bool = this.builtins.bool;
     this.set = this.builtins.set;
     this.tuple = this.builtins.tuple;
+    this.Ellipsis = this.builtins.Ellipsis;
 
     // Initialize arguments and executable path,
     // since some modules expect them to be set.
@@ -783,3 +804,59 @@ export class Python {
  * this object, such as `str`, `int`, `tuple`, etc.
  */
 export const python = new Python();
+
+/**
+ * Returns true if the value can be converted into a Python slice or
+ * slice tuple.
+ */
+function isSlice(value: unknown): boolean {
+  if (typeof value !== "string") return false;
+  if (!value.includes(":") && !value.includes("...")) return false;
+  return value
+    .split(",")
+    .map((item) => (
+      SliceItemRegExp.test(item) || // Slice
+      /^\s*-?\d+\s*$/.test(item) || // Number
+      /^\s*\.\.\.\s*$/.test(item) // Ellipsis
+    ))
+    .reduce((a, b) => a && b, true);
+}
+
+/**
+ * Returns a PyObject that is either a slice or a tuple of slices.
+ */
+function toSlice(sliceList: string): PyObject {
+  if (sliceList.includes(",")) {
+    const pySlicesHandle = sliceList.split(",")
+      .map(toSlice)
+      .map((pyObject) => pyObject.handle);
+
+    const pyTuple_Pack = new Deno.UnsafeFnPointer(py.PyTuple_Pack, {
+      parameters: ["i32", ...pySlicesHandle.map(() => "pointer" as const)],
+      result: "pointer",
+    });
+
+    const pyTupleHandle = pyTuple_Pack.call(
+      pySlicesHandle.length,
+      ...pySlicesHandle,
+    );
+    return new PyObject(pyTupleHandle);
+  } else if (/^\s*-?\d+\s*$/.test(sliceList)) {
+    return PyObject.from(parseInt(sliceList));
+  } else if (/^\s*\.\.\.\s*$/.test(sliceList)) {
+    return PyObject.from(python.Ellipsis);
+  } else {
+    const [start, stop, step] = sliceList
+      .split(":")
+      .map((
+        bound,
+      ) => (/^\s*-?\d+\s*$/.test(bound) ? parseInt(bound) : undefined));
+
+    const pySliceHandle = py.PySlice_New(
+      PyObject.from(start).handle,
+      PyObject.from(stop).handle,
+      PyObject.from(step).handle,
+    );
+    return new PyObject(pySliceHandle);
+  }
+}
diff --git a/src/symbols.ts b/src/symbols.ts
@@ -368,4 +368,8 @@ export const SYMBOLS = {
     parameters: ["pointer", "i32"],
     result: "pointer",
   },
+
+  PyTuple_Pack: {
+    type: "pointer",
+  },
 } as const;
diff --git a/src/util.ts b/src/util.ts
@@ -52,3 +52,11 @@ export function cstr(str: string): Uint8Array {
   encoder.encodeInto(str, buf);
   return buf;
 }
+
+/**
+ * Regular Expression used to test if a string is a `proper_slice`.
+ *
+ * Based on https://docs.python.org/3/reference/expressions.html#slicings
+ */
+export const SliceItemRegExp =
+  /^\s*(-?\d+)?\s*:\s*(-?\d+)?\s*(:\s*(-?\d+)?\s*)?$/;
diff --git a/test/test.ts b/test/test.ts
@@ -43,7 +43,13 @@ Deno.test("types", async (t) => {
 
   await t.step("dict", () => {
     const value = python.dict({ a: 1, b: 2 });
-    assertEquals(value.valueOf(), new Map([["a", 1], ["b", 2]]));
+    assertEquals(
+      value.valueOf(),
+      new Map([
+        ["a", 1],
+        ["b", 2],
+      ]),
+    );
   });
 
   await t.step("set", () => {
@@ -126,7 +132,9 @@ class Person:
 Deno.test("named argument", async (t) => {
   await t.step("single named argument", () => {
     assertEquals(
-      python.str("Hello, {name}!").format(new NamedArgument("name", "world"))
+      python
+        .str("Hello, {name}!")
+        .format(new NamedArgument("name", "world"))
         .valueOf(),
       "Hello, world!",
     );
@@ -171,3 +179,87 @@ Deno.test("custom proxy", () => {
   // Then, we use the wrapped proxy as if it were an original PyObject
   assertEquals(np.add(arr, 2).tolist().valueOf(), [3, 4, 5]);
 });
+
+Deno.test("slice", async (t) => {
+  await t.step("get", () => {
+    const list = python.list([1, 2, 3, 4, 5, 6, 7, 8, 9]);
+    assertEquals(list["1:"].valueOf(), [2, 3, 4, 5, 6, 7, 8, 9]);
+    assertEquals(list["1:2"].valueOf(), [2]);
+    assertEquals(list[":2"].valueOf(), [1, 2]);
+    assertEquals(list[":2:"].valueOf(), [1, 2]);
+    assertEquals(list["0:3:2"].valueOf(), [1, 3]);
+    assertEquals(list["-2:"].valueOf(), [8, 9]);
+    assertEquals(list["::2"].valueOf(), [1, 3, 5, 7, 9]);
+  });
+
+  await t.step("set", () => {
+    const np = python.import("numpy");
+    let list = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]);
+    list["1:"] = -5;
+    assertEquals(list.tolist().valueOf(), [1, -5, -5, -5, -5, -5, -5, -5, -5]);
+
+    list = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]);
+    list["1::3"] = -5;
+    assertEquals(list.tolist().valueOf(), [1, -5, 3, 4, -5, 6, 7, -5, 9]);
+
+    list = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]);
+    list["1:2:3"] = -5;
+    assertEquals(list.tolist().valueOf(), [1, -5, 3, 4, 5, 6, 7, 8, 9]);
+  });
+});
+
+Deno.test("slice list", async (t) => {
+  const np = python.import("numpy");
+
+  await t.step("get", () => {
+    const array = np.array([
+      [1, 2, 3],
+      [4, 5, 6],
+      [7, 8, 9],
+    ]);
+    assertEquals(array["0, :"].tolist().valueOf(), [1, 2, 3]);
+    assertEquals(array["1:, ::2"].tolist().valueOf(), [
+      [4, 6],
+      [7, 9],
+    ]);
+    assertEquals(array["1:, 0"].tolist().valueOf(), [4, 7]);
+  });
+
+  await t.step("set", () => {
+    const array = np.arange(15).reshape(3, 5);
+    array["1:, ::2"] = -99;
+    assertEquals(array.tolist().valueOf(), [
+      [0, 1, 2, 3, 4],
+      [-99, 6, -99, 8, -99],
+      [-99, 11, -99, 13, -99],
+    ]);
+  });
+
+  await t.step("whitespaces", () => {
+    const array = np.array([
+      [1, 2, 3],
+      [4, 5, 6],
+      [7, 8, 9],
+    ]);
+    assertEquals(array[" 1  :  , : :  2         "].tolist().valueOf(), [
+      [4, 6],
+      [7, 9],
+    ]);
+  });
+
+  await t.step("3d slicing", () => {
+    const a3 = np.array([[[10, 11, 12], [13, 14, 15], [16, 17, 18]], [
+      [20, 21, 22],
+      [23, 24, 25],
+      [26, 27, 28],
+    ], [[30, 31, 32], [33, 34, 35], [36, 37, 38]]]);
+
+    assertEquals(a3["0, :, 1"].tolist().valueOf(), [11, 14, 17]);
+  });
+
+  await t.step("ellipsis", () => {
+    const a4 = np.arange(16).reshape(2, 2, 2, 2);
+
+    assertEquals(a4["1, ..., 1"].tolist().valueOf(), [[9, 11], [13, 15]]);
+  });
+});