diff --git a/lib/elixir/lib/enum.ex b/lib/elixir/lib/enum.ex index 12d21286fda..239cbb69d4a 100644 --- a/lib/elixir/lib/enum.ex +++ b/lib/elixir/lib/enum.ex @@ -3217,19 +3217,30 @@ defmodule Enum do iex> Enum.sum(1..10//2) 25 + iex> Enum.sum(1..10, & &1 * 2) + 209 + + iex> Enum.sum([%{num: 1}, %{num: 2}], & &1.num) + 3 + + iex> Enum.sum(%{num_2: 2, num_3: 3}, fn {_, v} -> v end) + 5 """ - @spec sum(t) :: number - def sum(enumerable) + @spec sum(t, (element -> number)) :: number + def sum(enumerable, fun \\ fn item -> item end) + + def sum(first..last//step, fun) do + first = fun.(first) + last = fun.(last) - def sum(first..last//step = range) do - range + first..last//step |> Range.size() |> Kernel.*(first + last - rem(last - first, step)) |> div(2) end - def sum(enumerable) do - reduce(enumerable, 0, &+/2) + def sum(enumerable, fun) do + reduce(enumerable, 0, &(fun.(&1) + &2)) end @doc """ diff --git a/lib/elixir/test/elixir/enum_test.exs b/lib/elixir/test/elixir/enum_test.exs index 7269defd553..3ea9e184005 100644 --- a/lib/elixir/test/elixir/enum_test.exs +++ b/lib/elixir/test/elixir/enum_test.exs @@ -1254,6 +1254,39 @@ defmodule EnumTest do end end + test "sum/2" do + assert Enum.sum( + [%{number: 1, not_number: "foo"}, %{number: 2, not_number: "bar"}], + & &1.number + ) == 3 + + assert Enum.sum([%{number: 1}, %{number: 2}], &(&1.number * 2)) == 6 + assert Enum.sum(%{num_1: 1, num_2: 2, num_3: 3}, fn {_, v} -> v end) == 6 + assert Enum.sum(42..42, &(&1 * 2)) == 84 + assert Enum.sum(42..42, &div(&1, 2)) == 21 + assert Enum.sum(11..17, &(&1 * 2)) == 364 + assert Enum.sum(11..17, &div(&1, 2)) == 26 + + assert_raise ArgumentError, fn -> + Enum.sum(11..17, &(&1 / 2)) + end + + assert_raise ArgumentError, fn -> + Enum.sum(42..42, &(&1 / 2)) + end + + assert_raise ArithmeticError, fn -> + Enum.sum( + [%{number: 1, not_number: "foo"}, %{number: 2, not_number: "bar"}], + & &1.not_number + ) + end + + assert_raise ArithmeticError, fn -> + Enum.sum(%{letter_a: "a", letter_b: "b", letter_c: "c"}, fn {_, v} -> v end) + end + end + test "product/1" do assert Enum.product([]) == 1 assert Enum.product([1]) == 1