Support numpy.argsort kind keyword argument

pythran/pythonic/include/numpy/argsort.hpp

@@ -10,11 +10,16 @@ namespace numpy
 {
   template <class E>
   types::ndarray<long, types::array<long, 1>> argsort(E const &expr,
-                                                      types::none_type);
+                                                      types::none_type,
+                                                      types::none_type={});
 
   template <class T, class pS>
   types::ndarray<long, pS> argsort(types::ndarray<T, pS> const &a,
-                                   long axis = -1);
+                                   long axis = -1,  types::none_type kind={});
+
+  template <class T, class pS>
+  types::ndarray<long, pS> argsort(types::ndarray<T, pS> const &a,
+                                   long axis,  types::str const& kind);
 
   NUMPY_EXPR_TO_NDARRAY0_DECL(argsort);
 

pythran/pythonic/numpy/argsort.hpp

@@ -2,26 +2,22 @@
 #define PYTHONIC_NUMPY_ARGSORT_HPP
 
 #include "pythonic/include/numpy/argsort.hpp"
-
-#include "pythonic/types/ndarray.hpp"
-#include "pythonic/utils/allocate.hpp"
-#include "pythonic/utils/functor.hpp"
-#include "pythonic/utils/pdqsort.hpp"
+#include "pythonic/numpy/ndarray/sort.hpp"
 
 PYTHONIC_NS_BEGIN
 
 namespace numpy
 {
   template <class E>
   types::ndarray<long, types::array<long, 1>> argsort(E const &expr,
-                                                      types::none_type)
+                                                      types::none_type, types::none_type)
   {
     auto out = functor::array{}(expr).flat();
     return argsort(out);
   }
 
-  template <class T, class pS>
-  types::ndarray<long, pS> argsort(types::ndarray<T, pS> const &a, long axis)
+  template <class T, class pS, class Sorter>
+  types::ndarray<long, pS> _argsort(types::ndarray<T, pS> const &a, long axis, Sorter sorter)
   {
     constexpr auto N = std::tuple_size<pS>::value;
     if (axis < 0)
@@ -39,7 +35,7 @@ namespace numpy
         // fill with the original indices
         std::iota(iter_indices, iter_indices + step, 0L);
         // sort the index using the value from a
-        pdqsort(iter_indices, iter_indices + step,
+        sorter(iter_indices, iter_indices + step,
                 [a_base](long i1, long i2) { return a_base[i1] < a_base[i2]; });
       }
     } else {
@@ -57,7 +53,7 @@ namespace numpy
       std::iota(buffer_start, buffer_end, 0L);
       for (long i = 0; i < n; i++) {
         auto a_base = a.fbegin() + ith;
-        pdqsort(buffer, buffer + buffer_size,
+        sorter(buffer, buffer + buffer_size,
                 [a_base, stepper](long i1, long i2) {
                   return a_base[i1 * stepper] < a_base[i2 * stepper];
                 });
@@ -75,6 +71,23 @@ namespace numpy
     return indices;
   }
 
+  template <class T, class pS>
+  types::ndarray<long, pS> argsort(types::ndarray<T, pS> const &a, long axis, types::none_type) {
+    return _argsort(a, axis, ndarray::quicksorter());
+  }
+
+  template <class T, class pS>
+  types::ndarray<long, pS> argsort(types::ndarray<T, pS> const &a, long axis, types::str const& kind)
+  {
+      if (kind == "mergesort")
+        return _argsort(a, axis, ndarray::mergesorter());
+      else if (kind == "heapsort")
+        return _argsort(a, axis, ndarray::heapsorter());
+      else if (kind == "stable")
+        return _argsort(a, axis, ndarray::stablesorter());
+      return _argsort(a, axis, ndarray::quicksorter());
+  }
+
   NUMPY_EXPR_TO_NDARRAY0_IMPL(argsort);
 } // namespace numpy
 PYTHONIC_NS_END

pythran/tests/test_numpy_func2.py

@@ -628,6 +628,12 @@ def test_argsort3(self):
     def test_argsort4(self):
         self.run_test("def np_argsort4(x): return x.argsort(axis=None)", numpy.array([[3, 1, 2], [1 , 2, 3]]), np_argsort4=[NDArray[int,:,:]])
 
+    def test_argsort5(self):
+        self.run_test("def np_argsort5(x): return x.argsort(axis=None,kind=None)", numpy.array([[3, 1, 2], [1 , 2, 3]]), np_argsort5=[NDArray[int,:,:]])
+
+    def test_argsort6(self):
+        self.run_test("def np_argsort6(x): return x.argsort(kind='stable')", numpy.array([[3, 1, 2], [1 , 2, 3]]), np_argsort6=[NDArray[int,:,:]])
+
     def test_argmax0(self):
         self.run_test("def np_argmax0(a): return a.argmax()", numpy.arange(6).reshape(2,3), np_argmax0=[NDArray[int,:,:]])