feat: check for ProxiedPyObject symbol when converting object (denosaurs#19)

tjosepo · web-flow · commit c7d6bd4a45da · 2022-04-27T13:13:40.000+05:30
diff --git a/src/python.ts b/src/python.ts
@@ -2,6 +2,24 @@
 import { py } from "./ffi.ts";
 import { cstr } from "./util.ts";
 
+/**
+ * Symbol used on proxied Python objects to point to the original PyObject object.
+ * Can be used to implement PythonProxy and create your own proxies.
+ *
+ * See `PyObject#proxy` for more info on proxies.
+ */
+export const ProxiedPyObject = Symbol("ProxiedPyObject");
+
+/**
+ * Proxied Python object.
+ *
+ * When an object implements this interface, the object will not be converted
+ * to a Python Object and the original PyObject will be used.
+ */
+export interface PythonProxy {
+  [ProxiedPyObject]: PyObject;
+}
+
 /**
  * JS types that can be converted to Python Objects.
  *
@@ -30,6 +48,8 @@ import { cstr } from "./util.ts";
  * - `Set` becomes `set` in Python.
  *
  * If you pass a PyObject, it is used as-is.
+ *
+ * If you pass a PythonProxy, its original PyObject will be used.
  */
 export type PythonConvertible =
   | number
@@ -41,18 +61,12 @@ export type PythonConvertible =
   | string
   // deno-lint-ignore ban-types
   | Symbol
+  | PythonProxy
   | PythonConvertible[]
   | { [key: string]: PythonConvertible }
   | Map<PythonConvertible, PythonConvertible>
   | Set<PythonConvertible>;
 
-/**
- * Symbol used on proxied Python objects to point to the original PyObject object.
- *
- * See `PyObject#proxy` for more info on proxies.
- */
-export const ProxiedPyObject = Symbol("ProxiedPyObject");
-
 /**
  * An argument that can be passed to PyObject calls to indicate that the
  * argument should be passed as a named one.
@@ -310,6 +324,9 @@ export class PyObject {
       case "object": {
         if (v === null) {
           return python.builtins.None[ProxiedPyObject];
+        } else if (ProxiedPyObject in v) {
+          const proxy = v as PythonProxy;
+          return proxy[ProxiedPyObject];
         } else if (Array.isArray(v)) {
           const list = py.PyList_New(v.length) as Deno.UnsafePointer;
           for (let i = 0; i < v.length; i++) {
diff --git a/test/test.ts b/test/test.ts
@@ -1,5 +1,11 @@
 import { assert, assertEquals } from "./deps.ts";
-import { NamedArgument, PyObject, python } from "../mod.ts";
+import {
+  NamedArgument,
+  ProxiedPyObject,
+  PyObject,
+  python,
+  PythonProxy,
+} from "../mod.ts";
 
 const { version, executable } = python.import("sys");
 console.log("Python version:", version);
@@ -146,3 +152,22 @@ class Test:
 Deno.test("numpy", () => {
   const _np = python.import("numpy");
 });
+
+Deno.test("custom proxy", () => {
+  const np = python.import("numpy");
+
+  // We declare our own PythonProxy wrapper
+  const CustomProxy = class implements PythonProxy {
+    public readonly [ProxiedPyObject]: PyObject;
+
+    constructor(array: PythonProxy) {
+      this[ProxiedPyObject] = array[ProxiedPyObject];
+    }
+  };
+
+  // Wrap the result in our custom wrapper
+  const arr = new CustomProxy(np.array([1, 2, 3]));
+
+  // Then, we use the wrapped proxy as if it were an original PyObject
+  assertEquals(np.add(arr, 2).tolist().valueOf(), [3, 4, 5]);
+});
