Add Stream.take_every/2

Author: José Valim

lib/elixir/lib/enum.ex

@@ -659,9 +659,8 @@ defmodule Enum do
   end
 
   def filter_map(collection, filter, mapper) do
-    reduce(collection, [], fn(entry, acc) ->
-      if filter.(entry), do: [mapper.(entry)|acc], else: acc
-    end) |> :lists.reverse
+    Enumerable.reduce(collection, { :cont, [] }, R.filter_map(filter, mapper))
+    |> elem(1) |> :lists.reverse
   end
 
   @doc """
@@ -1490,17 +1489,9 @@ defmodule Enum do
   @spec take_every(t, integer) :: list
   def take_every(_collection, 0), do: []
   def take_every(collection, nth) do
-    res =
-      reduce(collection, nil, fn
-        x, { acc, ^nth }  -> { [x | acc], 1 }
-        _, { acc, count } -> { acc, count + 1 }
-        x, nil            -> { [x], 1 }
-      end)
-
-    case res do
-      nil -> []
-      { list, _ } -> :lists.reverse(list)
-    end
+    { _, { res, _ } } =
+      Enumerable.reduce(collection, { :cont, { [], :first } }, R.take_every(nth))
+    :lists.reverse(res)
   end
 
   @doc """

lib/elixir/lib/stream.ex

@@ -208,7 +208,7 @@ defmodule Stream do
   end
 
   @doc """
-  Creates a stream that will filter elements according to
+  Creates a stream that filters elements according to
   the given function on enumeration.
 
   ## Examples
@@ -223,6 +223,24 @@ defmodule Stream do
     lazy enum, fn(f1) -> R.filter(fun, f1) end
   end
 
+  @doc """
+  Creates a stream that filters and then map elements according
+  to given funtions.
+
+  It exists for simmetry with Enum.filter_map/3.
+
+  ## Examples
+
+      iex> stream = Stream.filter_map([1, 2, 3], fn(x) -> rem(x, 2) == 0 end)
+      iex> Enum.to_list(stream)
+      [2]
+
+  """
+  @spec filter_map(Enumerable.t, (element -> as_boolean(term)), (element -> any)) :: Enumerable.t
+  def filter_map(enum, filter, mapper) do
+    lazy enum, fn(f1) -> R.filter_map(filter, mapper, f1) end
+  end
+
   @doc """
   Creates a stream that will apply the given function on
   enumeration.
@@ -330,7 +348,26 @@ defmodule Stream do
     lazy enum, n, fn(f1) -> R.take(f1) end
   end
 
-  def take(_enum, 0), do: Lazy[enum: [], funs: [&(&1)]]
+  def take(_enum, 0), do: Lazy[enum: []]
+
+  @doc """
+  Creates a stream that takes every `n` item from the enumerable.
+
+  The first item is always included, unless n is 0.
+
+  ## Examples
+
+      iex> stream = Stream.take_every(1..10, 2)
+      iex> Enum.to_list(stream)
+      [1,3,5,7,9]
+
+  """
+  @spec take(Enumerable.t, non_neg_integer) :: Enumerable.t
+  def take_every(enum, n) when n > 0 do
+    lazy enum, n, fn(f1) -> R.take_every(n, f1) end
+  end
+
+  def take_every(_enum, 0), do: Lazy[enum: []]
 
   @doc """
   Lazily takes elements of the enumerable while the given

lib/elixir/lib/stream/reducers.ex

@@ -37,6 +37,18 @@ defmodule Stream.Reducers do
     end
   end
 
+  defmacro filter_map(filter, mapper, f // nil) do
+    quote do
+      fn(entry, acc) ->
+        if unquote(filter).(entry) do
+          cont(unquote(f), unquote(mapper).(entry), acc)
+        else
+          { :cont, acc }
+        end
+      end
+    end
+  end
+
   defmacro map(callback, f // nil) do
     quote do
       fn(entry, acc) ->
@@ -69,6 +81,18 @@ defmodule Stream.Reducers do
     end
   end
 
+  defmacro take_every(nth, f // nil) do
+    quote do
+      fn
+        entry, acc(h, n, t) when n === :first
+                            when n === unquote(nth) ->
+          cont_with_acc(unquote(f), entry, h, 1, t)
+        entry, acc(h, n, t) ->
+          { :cont, acc(h, n+1, t) }
+      end
+    end
+  end
+
   defmacro take_while(callback, f // nil) do
     quote do
       fn(entry, acc) ->

lib/elixir/test/elixir/stream_test.exs

@@ -126,6 +126,16 @@ defmodule StreamTest do
     assert Stream.filter(nats, &(rem(&1, 2) == 0)) |> Enum.take(5) == [2,4,6,8,10]
   end
 
+  test "filter_map" do
+    stream = Stream.filter_map([1,2,3], fn(x) -> rem(x, 2) == 0 end, &(&1 * 2))
+    assert is_lazy(stream)
+    assert Enum.to_list(stream) == [4]
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Stream.filter_map(nats, &(rem(&1, 2) == 0), &(&1 * 2))
+           |> Enum.take(5) == [4,8,12,16,20]
+  end
+
   test "flat_map" do
     stream = Stream.flat_map([1, 2, 3], &[&1, &1 * 2])
     assert is_lazy(stream)
@@ -228,6 +238,18 @@ defmodule StreamTest do
     assert Enum.zip(list, list) == Enum.zip(stream, stream)
   end
 
+  test "take_every" do
+    assert 1..10
+           |> Stream.take_every(2)
+           |> Enum.to_list == [1, 3, 5, 7, 9]
+
+    assert 1..10
+           |> Stream.drop(2)
+           |> Stream.take_every(2)
+           |> Stream.drop(1)
+           |> Enum.to_list == [5, 7, 9]
+  end
+
   test "take_while" do
     stream = Stream.take_while(1..1000, &(&1 <= 5))
     assert is_lazy(stream)