Fixes to getBuffer (#1399)

packages/numpy/test_numpy.py

@@ -287,8 +287,34 @@ def test_get_buffer_roundtrip(selenium, arg):
     )
 
 
+def test_get_buffer_big_endian(selenium):
+    selenium.run_js(
+        """
+        window.a = await pyodide.runPythonAsync(`
+            import numpy as np
+            np.arange(24, dtype="int16").byteswap().newbyteorder()
+        `);
+        """
+    )
+    with pytest.raises(
+        Exception, match="Javascript has no native support for big endian buffers"
+    ):
+        selenium.run_js("a.getBuffer()")
+    result = selenium.run_js(
+        """
+        let buf = a.getBuffer("i8")
+        let result = Array.from(buf.data);
+        buf.release();
+        a.destroy();
+        return result;
+        """
+    )
+    assert len(result) == 48
+    assert result[:18] == [0, 0, 0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, 8]
+
+
 def test_get_buffer_error_messages(selenium):
-    with pytest.raises(Exception, match="Javascript has no Float16Array"):
+    with pytest.raises(Exception, match="Javascript has no Float16 support"):
         selenium.run_js(
             """
             await pyodide.runPythonAsync(`

src/core/pyproxy.js

@@ -669,34 +669,26 @@ TEMP_EMJS_HELPER(() => {0, /* Magic, see comment */
 
   Module.PyProxyCallableMethods = {prototype : Function.prototype};
 
+  // clang-format off
   let type_to_array_map = new Map([
     [ "i8", Int8Array ],
     [ "u8", Uint8Array ],
+    [ "u8clamped", Uint8ClampedArray ],
     [ "i16", Int16Array ],
     [ "u16", Uint16Array ],
     [ "i32", Int32Array ],
     [ "u32", Uint32Array ],
     [ "i32", Int32Array ],
     [ "u32", Uint32Array ],
+    // if these aren't available, will be globalThis.BigInt64Array will be
+    // undefined rather than raising a ReferenceError.
+    [ "i64", globalThis.BigInt64Array],
+    [ "u64", globalThis.BigUint64Array],
     [ "f32", Float32Array ],
     [ "f64", Float64Array ],
-    // Python type formats
-    [ "b", Int8Array ],
-    [ "B", Uint8Array ],
-    [ "h", Int16Array ],
-    [ "H", Uint16Array ],
-    [ "i", Int32Array ],
-    [ "I", Uint32Array ],
-    [ "f", Float32Array ],
-    [ "d", Float64Array ],
+    [ "dataview", DataView ],
   ]);
-
-  if (globalThis.BigInt64Array) {
-    type_to_array_map.set("i64", BigInt64Array);
-    type_to_array_map.set("u64", BigUint64Array);
-    type_to_array_map.set("q", BigInt64Array);
-    type_to_array_map.set("Q", BigUint64Array);
-  }
+  // clang-format on
 
   Module.PyProxyBufferMethods = {
     /**
@@ -712,14 +704,14 @@ TEMP_EMJS_HELPER(() => {0, /* Magic, see comment */
      * suboffets (using e.g., ``np.ascontiguousarray``).
      *
      * @param {string} type The type of the desired output. Should be one of:
-     *    "i8", "u8", "i16", "u16", "i32", "u32", "i32", "u32", "i64", "u64",
-     *    "f32", or "f64,
+     *    "i8", "u8", "u8clamped", "i16", "u16", "i32", "u32", "i32", "u32",
+     *    "i64", "u64", "f32", "f64, or "dataview".
      * @returns PyBuffer
      */
-    getBuffer : function(type = "u8") {
+    getBuffer : function(type) {
       let ArrayType = undefined;
       if (type) {
-        let ArrayType = type_to_array_map.get(type);
+        ArrayType = type_to_array_map.get(type);
         if (ArrayType === undefined) {
           throw new Error(`Unknown type ${type}`);
         }
@@ -747,38 +739,27 @@ TEMP_EMJS_HELPER(() => {0, /* Magic, see comment */
       let c_contiguous = !!HEAP32[cur_ptr++];
       let f_contiguous = !!HEAP32[cur_ptr++];
 
-      _PyMem_Free(buffer_struct_ptr);
-
       let format = UTF8ToString(format_ptr);
+      _PyMem_Free(buffer_struct_ptr);
 
       let success = false;
       try {
+        let bigEndian = false;
         if (ArrayType === undefined) {
-          // Try to determine correct type from format.
-          // To understand this code it will be helpful to look at the tables
-          // here: https://docs.python.org/3/library/struct.html#format-strings
-          if (format.includes("e")) {
-            throw new Error("Javascript has no Float16Array.");
-          }
-          let cleaned_format = format;
-          // Normalize same-sized types
-          cleaned_format = cleaned_format.replace(/[spc?]/g, "B");
-          cleaned_format = cleaned_format.replace(/[nl]/g, "i");
-          cleaned_format = cleaned_format.replace(/[NLP]/g, "I");
-          let type_char = cleaned_format[0];
-          ArrayType = type_to_array_map.get(type_char);
-          if (ArrayType === undefined) {
-            if (/[qQ]/.test(type_char)) {
-              throw new Error(
-                  "64 bit integer formats (q and Q) are not supported in browsers without BigInt support. You must pass a type argument.");
-            } else {
-              throw new Error(
-                  "Unrecognized buffer format. You must pass a type argument.");
-            }
-          }
+          [ArrayType, bigEndian] = Module.processBufferFormatString(
+              format, " In this case, you can pass an explicit type argument.");
+        }
+        let alignment =
+            parseInt(ArrayType.name.replace(/[^0-9]/g, "")) / 8 || 1;
+        if (bigEndian && alignment > 1) {
+          throw new Error(
+              "Javascript has no native support for big endian buffers. " +
+              "In this case, you can pass an explicit type argument. " +
+              "For instance, `getBuffer('dataview')` will return a `DataView`" +
+              "which has native support for reading big endian data." +
+              "Alternatively, toJs will automatically convert the buffer " +
+              "to little endian.");
         }
-
-        let alignment = parseInt(ArrayType.name.replace(/[^0-9]/g, "")) / 8;
         if (startByteOffset % alignment !== 0 ||
             minByteOffset % alignment !== 0 ||
             maxByteOffset % alignment !== 0) {

src/tests/test_pyproxy.py

@@ -201,7 +201,7 @@ def test_pyproxy_iter(selenium):
 def test_pyproxy_get_buffer(selenium):
     selenium.run_js(
         """
-        await pyodide.runPython(`
+        pyodide.runPython(`
             from sys import getrefcount
             z1 = memoryview(bytes(range(24))).cast("b", [8,3])
             z2 = z1[-1::-1]
@@ -232,6 +232,58 @@ def test_pyproxy_get_buffer(selenium):
     )
 
 
+@pytest.mark.parametrize(
+    "array_type",
+    [
+        ["i8", "Int8Array", "b"],
+        ["u8", "Uint8Array", "B"],
+        ["u8clamped", "Uint8ClampedArray", "B"],
+        ["i16", "Int16Array", "h"],
+        ["u16", "Uint16Array", "H"],
+        ["i32", "Int32Array", "i"],
+        ["u32", "Uint32Array", "I"],
+        ["i64", "BigInt64Array", "q"],
+        ["u64", "BigUint64Array", "Q"],
+        ["f32", "Float32Array", "f"],
+        ["f64", "Float64Array", "d"],
+    ],
+)
+def test_pyproxy_get_buffer_type_argument(selenium, array_type):
+    selenium.run_js(
+        """
+        window.a = pyodide.runPython("bytes(range(256))");
+        """
+    )
+    try:
+        mv = memoryview(bytes(range(256)))
+        ty, array_ty, fmt = array_type
+        [check, result] = selenium.run_js(
+            f"""
+            let buf = a.getBuffer({ty!r});
+            let check = (buf.data.constructor.name === {array_ty!r});
+            let result = Array.from(buf.data);
+            if(typeof result[0] === "bigint"){{
+                result = result.map(x => x.toString(16));
+            }}
+            buf.release();
+            return [check, result];
+            """
+        )
+        assert check
+        if fmt.lower() == "q":
+            assert result == [hex(x).replace("0x", "") for x in list(mv.cast(fmt))]
+        elif fmt == "f" or fmt == "d":
+            from math import isclose
+
+            for a, b in zip(result, list(mv.cast(fmt))):
+                if a and b:
+                    assert isclose(a, b)
+        else:
+            assert result == list(mv.cast(fmt))
+    finally:
+        selenium.run_js("a.destroy(); window.a = undefined;")
+
+
 def test_pyproxy_mixins(selenium):
     result = selenium.run_js(
         """