Numpy sum imap

pythran/intrinsic.py

@@ -156,9 +156,10 @@ def isstaticfunction(self):
 
 class ConstMethodIntr(MethodIntr):
     def __init__(self, *combiners, **kwargs):
+        kwargs.setdefault('argument_effects',
+                          (ReadEffect(),) * 12)
         super(ConstMethodIntr, self).__init__(
-            *combiners,
-            argument_effects=(ReadEffect(),) * 12, **kwargs)
+            *combiners, **kwargs)
 
 
 class AttributeIntr(Intrinsic):

pythran/pythonic/include/__builtin__/xrange.hpp

@@ -51,6 +51,9 @@ namespace pythonic
       using const_iterator = xrange_iterator;
       using reverse_iterator = xrange_riterator;
       using const_reverse_iterator = xrange_riterator;
+      using dtype = value_type;
+      static constexpr bool is_vectorizable = false;
+      static constexpr size_t value = 1;
 
       long _begin;
       long _end;

pythran/pythonic/include/itertools/ifilter.hpp

@@ -59,6 +59,9 @@ namespace pythonic
 
         using value_type = typename List0::value_type;
         using iterator = ifilter_iterator<Operator, List0>;
+        using dtype = typename types::dtype_of<value_type>::type;
+        static constexpr bool is_vectorizable = false;
+        static constexpr size_t value = types::list<value_type>::value;
 
         iterator end_iter;
 

pythran/pythonic/include/itertools/imap.hpp

@@ -91,6 +91,9 @@ namespace pythonic
                     imap_iterator<Operator, Iters...> {
         using value_type = typename imap_res<Operator, Iters...>::type;
         using iterator = imap_iterator<Operator, Iters...>;
+        using dtype = typename types::dtype_of<value_type>::type;
+        static constexpr bool is_vectorizable = false;
+        static constexpr size_t value = types::list<value_type>::value;
 
         iterator end_iter;
 

pythran/pythonic/include/utils/numpy_traits.hpp

@@ -3,6 +3,25 @@
 
 namespace pythonic
 {
+
+  namespace itertools
+  {
+    namespace details
+    {
+      template <class F, class... V>
+      struct imap;
+      template <class F, class ... V>
+      struct izip;
+      template <class F, class List0>
+      struct ifilter;
+    }
+  }
+
+  namespace __builtin__
+  {
+    struct xrange;
+  }
+
   namespace types
   {
 
@@ -118,6 +137,22 @@ namespace pythonic
       static constexpr bool value = true;
     };
 
+    template <class F, class... V>
+    struct is_numexpr_arg<itertools::details::imap<F, V...>> : std::true_type {
+    };
+
+    template <class F, class V>
+    struct is_numexpr_arg<itertools::details::ifilter<F, V>> : std::true_type {
+    };
+
+    template <class F, class... V>
+    struct is_numexpr_arg<itertools::details::izip<F, V...>> : std::true_type {
+    };
+
+    template <>
+    struct is_numexpr_arg<__builtin__::xrange> : std::true_type {
+    };
+
     template <class E>
     struct dtype_of {
       template <class T>

pythran/pythonic/itertools/ifilter.hpp

@@ -95,7 +95,8 @@ namespace pythonic
       template <typename Operator, typename List0>
       ifilter<Operator, List0>::ifilter(Operator _op, List0 const &_seq)
           : utils::iterator_reminder<false, List0>(_seq),
-            iterator(_op, this->value), end_iter(npos(), _op, this->value)
+            iterator(_op, utils::iterator_reminder<false, List0>::value),
+            end_iter(npos(), _op, utils::iterator_reminder<false, List0>::value)
       {
       }
 

pythran/pythonic/itertools/imap.hpp

@@ -180,9 +180,10 @@ namespace pythonic
           : utils::iterator_reminder<true, Iters...>(
                 std::forward<Types>(_iters)...),
             imap_iterator<Operator, Iters...>(
-                _op, this->value,
+                _op, utils::iterator_reminder<true, Iters...>::value,
                 typename utils::gens<sizeof...(Iters)>::type{}),
-            end_iter(npos(), _op, this->value,
+            end_iter(npos(), _op,
+                     utils::iterator_reminder<true, Iters...>::value,
                      typename utils::gens<sizeof...(Iters)>::type{})
       {
       }

pythran/tables.py

@@ -13,7 +13,7 @@
 from pythran.intrinsic import ClassWithReadOnceConstructor
 from pythran.intrinsic import ConstFunctionIntr, FunctionIntr, UpdateEffect
 from pythran.intrinsic import ConstMethodIntr, MethodIntr, AttributeIntr
-from pythran.intrinsic import ReadEffect, ConstantIntr
+from pythran.intrinsic import ReadEffect, ConstantIntr, ReadOnceEffect
 from pythran.intrinsic import ReadOnceFunctionIntr, ConstExceptionIntr
 from pythran.types.conversion import PYTYPE_TO_CTYPE_TABLE
 from pythran import range as prange
@@ -626,7 +626,8 @@ def update_effects(self, node):
         "std_": ConstMethodIntr(args=('a', 'axis', 'dtype'),
                                 defaults=(None, None)),
         "subtract": ConstFunctionIntr(),
-        "sum": ConstMethodIntr(),
+        "sum": ConstMethodIntr(argument_effects=(ReadOnceEffect(),
+                                                 ReadOnceEffect())),
         "swapaxes": ConstMethodIntr(),
         "take": ConstMethodIntr(),
         "tan": ConstFunctionIntr(),

pythran/tests/test_numpy_func1.py

@@ -42,6 +42,37 @@ def test_sum10_(self):
     def test_sum11_(self):
         self.run_test("def np_sum11_(a): import numpy as np ; return np.sum(a+a,2)", numpy.arange(12).reshape(2,3,2), np_sum11_=[numpy.array([[[int]]])])
 
+    def test_sum_imap(self):
+        self.run_test("""
+            def np_sum_imap(a):
+                import numpy as np
+                from itertools import imap
+                return np.sum(imap(lambda x: x + 2, a))
+                      """, range(10), np_sum_imap=[[int]])
+
+    def test_sum_ifilter(self):
+        self.run_test("""
+            def np_sum_ifilter(a):
+                import numpy as np
+                from itertools import ifilter
+                return np.sum(ifilter(lambda x: x % 2, a))
+                      """, range(10), np_sum_ifilter=[[int]])
+
+    def test_sum_izip(self):
+        self.run_test("""
+            def np_sum_izip(a):
+                import numpy as np
+                from itertools import izip
+                return np.sum(izip(a, a))
+                      """, range(10), np_sum_izip=[[int]])
+
+    def test_sum_xrange(self):
+        self.run_test("""
+            def np_sum_xrange(a):
+                import numpy as np
+                return np.sum(xrange(a))
+                      """, 10, np_sum_xrange=[int])
+
     def test_prod_(self):
         """ Check prod function for numpy array. """
         self.run_test("""