diff --git a/lib/elixir/lib/enum.ex b/lib/elixir/lib/enum.ex index 937860747c2..225a61a6a3e 100644 --- a/lib/elixir/lib/enum.ex +++ b/lib/elixir/lib/enum.ex @@ -2734,11 +2734,15 @@ defmodule Enum do [{1, :a}, {2, :b}, {3, :c}] """ - @spec zip([t]) :: t + @spec zip([t] | %Stream{}) :: t def zip([]), do: [] - def zip(enumerables) when is_list(enumerables) do + def zip(%{__struct__: Stream} = enumerables), do: do_zip(enumerables) + + def zip(enumerables) when is_list(enumerables), do: do_zip(enumerables) + + defp do_zip(enumerables) do Stream.zip(enumerables).({:cont, []}, &{:cont, [&1 | &2]}) |> elem(1) |> :lists.reverse() diff --git a/lib/elixir/lib/stream.ex b/lib/elixir/lib/stream.ex index 60edd5141d2..fdbb59b4d34 100644 --- a/lib/elixir/lib/stream.ex +++ b/lib/elixir/lib/stream.ex @@ -1106,8 +1106,12 @@ defmodule Stream do [{1, :a, "foo"}, {2, :b, "bar"}, {3, :c, "baz"}] """ - @spec zip([Enumerable.t()]) :: Enumerable.t() - def zip(enumerables) when is_list(enumerables) do + @spec zip([Enumerable.t()] | %Stream{}) :: Enumerable.t() + def zip(%Stream{} = enumerables), do: zip_fun(enumerables) + + def zip(enumerables) when is_list(enumerables), do: zip_fun(enumerables) + + defp zip_fun(enumerables) do step = &do_zip_step(&1, &2) enum_funs = diff --git a/lib/elixir/test/elixir/stream_test.exs b/lib/elixir/test/elixir/stream_test.exs index b894da30ac8..735929382a7 100644 --- a/lib/elixir/test/elixir/stream_test.exs +++ b/lib/elixir/test/elixir/stream_test.exs @@ -1084,6 +1084,10 @@ defmodule StreamTest do assert Stream.zip([concat, cycle]) |> Enum.to_list() == [{1, :a}, {2, :b}, {3, :c}, {4, :a}, {5, :b}, {6, :c}] + assert Stream.chunk_every([0, 1, 2, 3], 2) + |> Stream.zip() + |> Enum.to_list() == [{0, 2}, {1, 3}] + assert_raise FunctionClauseError, fn -> enum_of_enums = Stream.cycle([[1, 2], [:a, :b]]) Stream.zip(enum_of_enums)