Add step to Date.Range (#10816)

Author: wojtekmach

lib/elixir/lib/calendar/date.ex

@@ -83,28 +83,70 @@ defmodule Date do
       366
       iex> Enum.member?(range, ~D[2001-02-01])
       true
-      iex> Enum.reduce(range, 0, fn _date, acc -> acc - 1 end)
-      -366
+      iex> Enum.take(range, 3)
+      [~D[2001-01-01], ~D[2001-01-02], ~D[2001-01-03]]
 
   """
   @doc since: "1.5.0"
   @spec range(Calendar.date(), Calendar.date()) :: Date.Range.t()
   def range(%{calendar: calendar} = first, %{calendar: calendar} = last) do
     {first_days, _} = to_iso_days(first)
     {last_days, _} = to_iso_days(last)
+    # TODO: Deprecate inferring a range with step of -1 on Elixir v1.16
+    step = if first_days <= last_days, do: 1, else: -1
+    range(first, first_days, last, last_days, calendar, step)
+  end
+
+  def range(%{calendar: _, year: _, month: _, day: _}, %{calendar: _, year: _, month: _, day: _}) do
+    raise ArgumentError, "both dates must have matching calendars"
+  end
+
+  @doc """
+  Returns a range of dates with step.
+
+  ## Examples
+
+      iex> range = Date.range(~D[2001-01-01], ~D[2002-01-01], 2)
+      iex> range
+      #DateRange<~D[2001-01-01], ~D[2002-01-01], 2>
+      iex> Enum.count(range)
+      183
+      iex> Enum.member?(range, ~D[2001-01-03])
+      true
+      iex> Enum.take(range, 3)
+      [~D[2001-01-01], ~D[2001-01-03], ~D[2001-01-05]]
 
+  """
+  @doc since: "1.12.0"
+  @spec range(Calendar.date(), Calendar.date(), step :: pos_integer | neg_integer) ::
+          Date.Range.t()
+  def range(%{calendar: calendar} = first, %{calendar: calendar} = last, step)
+      when is_integer(step) and step != 0 do
+    {first_days, _} = to_iso_days(first)
+    {last_days, _} = to_iso_days(last)
+    range(first, first_days, last, last_days, calendar, step)
+  end
+
+  def range(
+        %{calendar: _, year: _, month: _, day: _} = first,
+        %{calendar: _, year: _, month: _, day: _} = last,
+        step
+      ) do
+    raise ArgumentError,
+          "both dates must have matching calendar and the step must be an integer " <>
+            "different than zero, got: #{inspect(first)}, #{inspect(last)}, #{step}"
+  end
+
+  defp range(first, first_days, last, last_days, calendar, step) do
     %Date.Range{
       first: %Date{calendar: calendar, year: first.year, month: first.month, day: first.day},
       last: %Date{calendar: calendar, year: last.year, month: last.month, day: last.day},
       first_in_iso_days: first_days,
-      last_in_iso_days: last_days
+      last_in_iso_days: last_days,
+      step: step
     }
   end
 
-  def range(%{calendar: _, year: _, month: _, day: _}, %{calendar: _, year: _, month: _, day: _}) do
-    raise ArgumentError, "both dates must have matching calendars"
-  end
-
   @doc """
   Returns the current date in UTC.
 
@@ -659,11 +701,12 @@ defmodule Date do
     end
   end
 
-  defp to_iso_days(%{calendar: Calendar.ISO, year: year, month: month, day: day}) do
+  @doc false
+  def to_iso_days(%{calendar: Calendar.ISO, year: year, month: month, day: day}) do
     {Calendar.ISO.date_to_iso_days(year, month, day), {0, 86_400_000_000}}
   end
 
-  defp to_iso_days(%{calendar: calendar, year: year, month: month, day: day}) do
+  def to_iso_days(%{calendar: calendar, year: year, month: month, day: day}) do
     calendar.naive_datetime_to_iso_days(year, month, day, 0, 0, 0, {0, 0})
   end
 

lib/elixir/lib/calendar/date_range.ex

@@ -2,12 +2,13 @@ defmodule Date.Range do
   @moduledoc """
   Returns an inclusive range between dates.
 
-  Ranges must be created with the `Date.range/2` function.
+  Ranges must be created with the `Date.range/2` or `Date.range/3` function.
 
   The following fields are public:
 
     * `:first` - the initial date on the range
     * `:last` - the last date on the range
+    * `:step` - (since v1.12.0) the step
 
   The remaining fields are private and should not be accessed.
   """
@@ -16,98 +17,98 @@ defmodule Date.Range do
           first: Date.t(),
           last: Date.t(),
           first_in_iso_days: iso_days(),
-          last_in_iso_days: iso_days()
+          last_in_iso_days: iso_days(),
+          step: pos_integer | neg_integer
         }
 
   @typep iso_days() :: Calendar.iso_days()
 
-  defstruct [:first, :last, :first_in_iso_days, :last_in_iso_days]
+  defstruct [:first, :last, :first_in_iso_days, :last_in_iso_days, :step]
 
   defimpl Enumerable do
     def member?(%{first: %{calendar: calendar}} = range, %Date{calendar: calendar} = date) do
       %{
-        first: first,
-        last: last,
-        first_in_iso_days: first_in_iso_days,
-        last_in_iso_days: last_in_iso_days
+        first_in_iso_days: first_days,
+        last_in_iso_days: last_days,
+        step: step
       } = range
 
-      %{year: first_year, month: first_month, day: first_day} = first
-      %{year: last_year, month: last_month, day: last_day} = last
-      %{year: year, month: month, day: day} = date
-      first = {first_year, first_month, first_day}
-      last = {last_year, last_month, last_day}
-      date = {year, month, day}
-
-      if first_in_iso_days <= last_in_iso_days do
-        {:ok, date >= first and date <= last}
-      else
-        {:ok, date >= last and date <= first}
+      {days, _} = Date.to_iso_days(date)
+
+      cond do
+        empty?(range) ->
+          {:ok, false}
+
+        first_days <= last_days ->
+          {:ok, first_days <= days and days <= last_days and rem(days - first_days, step) == 0}
+
+        true ->
+          {:ok, last_days <= days and days <= first_days and rem(days - first_days, step) == 0}
       end
     end
 
     def member?(_, _) do
       {:ok, false}
     end
 
-    def count(%{first_in_iso_days: first, last_in_iso_days: last}) do
-      {:ok, abs(first - last) + 1}
+    def count(range) do
+      {:ok, size(range)}
     end
 
     def slice(range) do
       %{
         first_in_iso_days: first,
-        last_in_iso_days: last,
-        first: %{calendar: calendar}
+        first: %{calendar: calendar},
+        step: step
       } = range
 
-      if first <= last do
-        {:ok, last - first + 1, &slice_asc(first + &1, &2, calendar)}
-      else
-        {:ok, first - last + 1, &slice_desc(first - &1, &2, calendar)}
-      end
+      {:ok, size(range), &slice(first + &1 * step, step, &2, calendar)}
     end
 
-    defp slice_asc(current, 1, calendar), do: [date_from_iso_days(current, calendar)]
-
-    defp slice_asc(current, remaining, calendar) do
-      [date_from_iso_days(current, calendar) | slice_asc(current + 1, remaining - 1, calendar)]
+    defp slice(current, _step, 1, calendar) do
+      [date_from_iso_days(current, calendar)]
     end
 
-    defp slice_desc(current, 1, calendar), do: [date_from_iso_days(current, calendar)]
-
-    defp slice_desc(current, remaining, calendar) do
-      [date_from_iso_days(current, calendar) | slice_desc(current - 1, remaining - 1, calendar)]
+    defp slice(current, step, remaining, calendar) do
+      [
+        date_from_iso_days(current, calendar)
+        | slice(current + step, step, remaining - 1, calendar)
+      ]
     end
 
     def reduce(range, acc, fun) do
       %{
-        first_in_iso_days: first_in_iso_days,
-        last_in_iso_days: last_in_iso_days,
-        first: %{calendar: calendar}
+        first_in_iso_days: first_days,
+        last_in_iso_days: last_days,
+        first: %{calendar: calendar},
+        step: step
       } = range
 
-      up? = first_in_iso_days <= last_in_iso_days
-      reduce(first_in_iso_days, last_in_iso_days, acc, fun, calendar, up?)
+      reduce(first_days, last_days, acc, fun, step, calendar)
     end
 
-    defp reduce(_x, _y, {:halt, acc}, _fun, _calendar, _up?) do
+    defp reduce(_first_days, _last_days, {:halt, acc}, _fun, _step, _calendar) do
       {:halted, acc}
     end
 
-    defp reduce(x, y, {:suspend, acc}, fun, calendar, up?) do
-      {:suspended, acc, &reduce(x, y, &1, fun, calendar, up?)}
+    defp reduce(first_days, last_days, {:suspend, acc}, fun, step, calendar) do
+      {:suspended, acc, &reduce(first_days, last_days, &1, fun, step, calendar)}
     end
 
-    defp reduce(x, y, {:cont, acc}, fun, calendar, up? = true) when x <= y do
-      reduce(x + 1, y, fun.(date_from_iso_days(x, calendar), acc), fun, calendar, up?)
+    defp reduce(first_days, last_days, {:cont, acc}, fun, step, calendar)
+         when step > 0 and first_days <= last_days
+         when step < 0 and first_days >= last_days do
+      reduce(
+        first_days + step,
+        last_days,
+        fun.(date_from_iso_days(first_days, calendar), acc),
+        fun,
+        step,
+        calendar
+      )
     end
 
-    defp reduce(x, y, {:cont, acc}, fun, calendar, up? = false) when x >= y do
-      reduce(x - 1, y, fun.(date_from_iso_days(x, calendar), acc), fun, calendar, up?)
-    end
-
-    defp reduce(_, _, {:cont, acc}, _fun, _calendar, _up) do
+    defp reduce(_, _, {:cont, acc}, _fun, _step, _calendar) do
       {:done, acc}
     end
 
@@ -122,11 +123,44 @@ defmodule Date.Range do
 
       %Date{year: year, month: month, day: day, calendar: calendar}
     end
+
+    defp size(%Date.Range{first_in_iso_days: first_days, last_in_iso_days: last_days, step: step})
+         when step > 0 and first_days > last_days,
+         do: 0
+
+    defp size(%Date.Range{first_in_iso_days: first_days, last_in_iso_days: last_days, step: step})
+         when step < 0 and first_days < last_days,
+         do: 0
+
+    defp size(%Date.Range{first_in_iso_days: first_days, last_in_iso_days: last_days, step: step}),
+      do: abs(div(last_days - first_days, step)) + 1
+
+    defp empty?(%Date.Range{
+           first_in_iso_days: first_days,
+           last_in_iso_days: last_days,
+           step: step
+         })
+         when step > 0 and first_days > last_days,
+         do: true
+
+    defp empty?(%Date.Range{
+           first_in_iso_days: first_days,
+           last_in_iso_days: last_days,
+           step: step
+         })
+         when step < 0 and first_days < last_days,
+         do: true
+
+    defp empty?(%Date.Range{}), do: false
   end
 
   defimpl Inspect do
-    def inspect(%Date.Range{first: first, last: last}, _) do
+    def inspect(%Date.Range{first: first, last: last, step: 1}, _) do
       "#DateRange<" <> inspect(first) <> ", " <> inspect(last) <> ">"
     end
+
+    def inspect(%Date.Range{first: first, last: last, step: step}, _) do
+      "#DateRange<" <> inspect(first) <> ", " <> inspect(last) <> ", #{step}>"
+    end
   end
 end