Add Numpy.to_bigarray_k

CHANGES.md

@@ -1,5 +1,13 @@
 [*] marks changes that break compatibility with previous versions.
 
+# Next version
+
+- `Numpy.to_bigarray_k` is continuation-passing-style version of `Numpy.to_bigarray`,
+  allowing caller to convert Numpy arrays to bigarrays without having to know
+  the kind and the layout of the array
+  (suggested by Lindsay Errington and Andie Sigler,
+  https://github.com/thierry-martinez/pyml/issues/81)
+
 # 2022-03-25
 
 - Fix debug build detection

numpy.ml

@@ -81,3 +81,33 @@ let to_bigarray kind layout t =
       "Numpy.to_bigarray: Numpy array has %s layout, but to_bigarray expected %s"
       (string_of_layout layout') (string_of_layout layout));
   array
+
+type ('a, 'b, 'c) to_bigarray =
+  { kind : ('a, 'b) Bigarray.kind
+  ; layout : 'c Bigarray.layout
+  ; array : ('a, 'b, 'c) Bigarray.Genarray.t
+  }
+
+type 'r to_bigarray_k =
+  { f : 'a 'b 'c . ('a, 'b, 'c) to_bigarray -> 'r }
+
+let to_bigarray_k (k : 'r to_bigarray_k) t : 'r =
+  if not (Py.Object.is_instance t (Py.Array.pyarray_type ())) then
+    invalid_arg "Numpy.to_bigarray";
+  let kind, layout, array = bigarray_of_pyarray (Py.Array.numpy_api ()) t in
+  k.f { kind; layout; array }
+
+external compare_kind :
+  ('a, 'b) Bigarray.kind -> ('c, 'd) Bigarray.kind -> int = "%compare"
+
+external compare_layout :
+  'a Bigarray.layout -> 'b Bigarray.layout -> int = "%compare"
+
+let check_kind_and_layout (kind : ('a, 'b) Bigarray.kind)
+      (layout : 'c Bigarray.layout) t :
+      ('a, 'b, 'c) Bigarray.Genarray.t option =
+  if compare_kind kind (Bigarray.Genarray.kind t) = 0 &&
+       compare_layout layout (Bigarray.Genarray.layout t) = 0 then
+    Some (Obj.magic t)
+  else
+    None

numpy.mli

@@ -31,4 +31,41 @@ val to_bigarray:
   ('a, 'b) Bigarray.kind -> 'c Bigarray.layout -> Py.Object.t ->
     ('a, 'b, 'c) Bigarray.Genarray.t
 (** [to_bigarray kind layout a] returns a bigarray that shares the same
-    contents than the Numpy array [a]. *)
+    contents than the Numpy array [a].
+    If `kind` and/or `layout` are unknown, you may use {!val:to_bigarray_k}. *)
+
+type ('a, 'b, 'c) to_bigarray =
+  { kind : ('a, 'b) Bigarray.kind
+  ; layout : 'c Bigarray.layout
+  ; array : ('a, 'b, 'c) Bigarray.Genarray.t
+  }
+
+type 'r to_bigarray_k =
+  { f : 'a 'b 'c . ('a, 'b, 'c) to_bigarray -> 'r }
+
+val to_bigarray_k : 'r to_bigarray_k -> Py.Object.t -> 'r
+(** [to_bigarray_k k a] calls [k.f] with the contents of the Numpy array [a].
+    [k.f] has to be polymorphic in the kind and the layout of the bigarray:
+    functions {!val:compare_kind}, {!val:compare_layout} and
+    {!val:check_kind_and_layout} can be used to introspect the bigarray
+    polymorphically. *)
+
+val compare_kind : ('a, 'b) Bigarray.kind -> ('c, 'd) Bigarray.kind -> int
+(** [compare_kind] provides a total order on {!val:Bigarray.kind}.
+    As opposed to generic [compare] of OCaml standard libary,
+    [compare_kind] is polymorphic in the kind of the bigarray. *)
+
+val compare_layout : 'a Bigarray.layout -> 'b Bigarray.layout -> int
+(** [compare_layout] provides a total order on {!val:Bigarray.layout}.
+    As opposed to generic [compare] of OCaml standard libary,
+    [compare_kind] is polymorphic in the layout of the bigarray. *)
+
+val check_kind_and_layout :
+  ('a, 'b) Bigarray.kind -> 'c Bigarray.layout ->
+    ('d, 'e, 'f) Bigarray.Genarray.t ->
+    ('a, 'b, 'c) Bigarray.Genarray.t option
+(** [check_kind_and_layout kind layout a] returns [Some a] if [a] has the given
+    [kind] and [layout] (that is to say, if we have the following type
+    equalities, ['a = 'd], ['b = 'e] and ['c = 'f]).
+    This function allows the callback of {!val:to_bigarray_k} to be polymorphic
+    in the kind of the array. *)

numpy_tests.ml

@@ -153,6 +153,38 @@ let () =
           Pyml_tests_common.Passed
         end)
 
+let () =
+  Pyml_tests_common.add_test ~title:"to_bigarray_k"
+    (fun () ->
+      if Py.Import.try_import_module "numpy" = None then
+        Pyml_tests_common.Disabled "numpy is not available"
+      else
+        begin
+          let m = Py.Import.add_module "test" in
+          let callback arg =
+            let k { Numpy.kind; layout; array } =
+              assert (Numpy.compare_kind kind Bigarray.nativeint = 0);
+              assert (Numpy.compare_layout layout Bigarray.c_layout = 0);
+              let bigarray =
+                Stdcompat.Option.get (Numpy.check_kind_and_layout
+                  Bigarray.nativeint Bigarray.c_layout array) in
+              assert (Bigarray.Genarray.dims bigarray = [| 4 |]);
+              let array1 = Bigarray.array1_of_genarray bigarray in
+              assert (Bigarray.Array1.get array1 0 = 0n);
+              assert (Bigarray.Array1.get array1 1 = 1n);
+              assert (Bigarray.Array1.get array1 2 = 2n);
+              assert (Bigarray.Array1.get array1 3 = 3n) in
+            Numpy.to_bigarray_k { Numpy.f = k } arg.(0);
+            Py.none in
+          Py.Module.set m "callback" (Py.Callable.of_function callback);
+          assert (Py.Run.simple_string "
+from test import callback
+import numpy
+callback(numpy.array([0,1,2,3]))
+");
+          Pyml_tests_common.Passed
+        end)
+
 let () =
   if not !Sys.interactive then
     Pyml_tests_common.main ()