Add Stream.after/2, Stream.run/2 and File.stream_to!/3 for side effects

Closes #1831

Author: José Valim

CHANGELOG.md

@@ -2,12 +2,13 @@
 
 * Enhancements
   * [Exception] Allow `exception/1` to be overriden and promote it as the main mechanism to customize exceptions
+  * [File] Add `File.stream_to/3`
   * [Kernel] Add `List.delete_at/2` and `List.updated_at/3`
   * [Kernel] Add `Enum.reverse/2`
   * [Kernel] Implement `defmodule/2`, `@/1`, `def/2` and friends in Elixir itself. `case/2`, `try/2` and `receive/1` have been made special forms. `var!/1`, `var!/2` and `alias!/1` have also been implemented in Elixir and demoted from special forms
   * [Record] Support dynamic fields in `defrecordp`
   * [Stream] Add `Stream.resource/3`
-  * [Stream] Add `Stream.zip/2`, `Stream.filter_map/3`, `Stream.each/2`, `Stream.take_every/2`, `Stream.chunks/2`, `Stream.chunks/3`, `Stream.chunks/4`, `Stream.chunks_by/2`, `Stream.scan/2`, `Stream.scan/3` and `Stream.uniq/2`
+  * [Stream] Add `Stream.zip/2`, `Stream.filter_map/3`, `Stream.each/2`, `Stream.take_every/2`, `Stream.chunks/2`, `Stream.chunks/3`, `Stream.chunks/4`, `Stream.chunks_by/2`, `Stream.scan/2`, `Stream.scan/3`, `Stream.uniq/2`, `Stream.after/2` and `Stream.run/1`
   * [Stream] Support `Stream.take/2` and `Stream.drop/2` with negative counts
   * [Typespec] Support `is_var/1` in typespecs
 

lib/elixir/lib/file.ex

@@ -954,12 +954,12 @@ defmodule File do
   end
 
   @doc """
-  Opens the given `file` with the given `mode` and returns
-  a stream for each `:line` (default) or for a given number
+  Opens the file at the given `path` with the given `modes` and
+  returns a stream for each `:line` (default) or for a given number
   of bytes given by `line_or_bytes`.
 
   The returned stream will fail for the same reasons as
-  `File.open!`. Note that the file is opened only and every time
+  `File.open!/2`. Note that the file is opened only and every time
   streaming begins.
 
   Note that stream by default uses `IO.binread/2` unless
@@ -988,6 +988,42 @@ defmodule File do
     Stream.resource(start_fun, next_fun, &F.close/1)
   end
 
+  @doc """
+  Receives a stream and returns a new stream that will open the file
+  at the given `path` for write  with the extra `modes` and write
+  each value to the file.
+
+  The returned stream will fail for the same reasons as
+  `File.open!/2`. Note that the file is opened only and every time
+  streaming begins.
+
+  Note that stream by default uses `IO.binwrite/2` unless
+  the file is opened with an encoding, then the slower `IO.write/2`
+  is used to do the proper data conversion and guarantees.
+  """
+  def stream_to!(stream, path, modes // []) do
+    modes = open_defaults([:write|List.delete(modes, :write)], true)
+    bin   = nil? List.keyfind(modes, :encoding, 0)
+
+    fn acc, f ->
+      case F.open(path, modes) do
+        { :ok, device } ->
+          each =
+            case bin do
+              true  -> &IO.binwrite(device, &1)
+              false -> &IO.write(device, &1)
+            end
+
+          stream
+          |> Stream.each(each)
+          |> Stream.after(fn -> F.close(device) end)
+          |> Enumerable.Stream.Lazy.reduce(acc, f)
+        { :error, reason } ->
+          raise File.Error, reason: reason, action: "stream_to", path: to_string(path)
+      end
+    end
+  end
+
   @doc false
   def binstream!(file, mode // [], line_or_bytes // :line) do
     IO.write "File.binstream! is deprecated, simply use File.stream! instead\n" <>

lib/elixir/lib/stream.ex

@@ -83,7 +83,7 @@ defmodule Stream do
   like `Stream.cycle/1`, `Stream.unfold/2`, `Stream.resource/3` and more.
   """
 
-  defrecord Lazy, enum: nil, funs: [], accs: [], done: nil
+  defrecord Lazy, enum: nil, funs: [], accs: [], after: [], last: nil
 
   defimpl Enumerable, for: Lazy do
     @compile :inline_list_funs
@@ -103,35 +103,32 @@ defmodule Stream do
       { :error, __MODULE__ }
     end
 
-    defp do_reduce(Lazy[enum: enum, funs: funs, accs: accs, done: done], acc, fun) do
+    defp do_reduce(Lazy[enum: enum, funs: funs, accs: accs, last: last, after: after_funs], acc, fun) do
       composed = :lists.foldl(fn fun, acc -> fun.(acc) end, fun, funs)
-      do_each(&Enumerable.reduce(enum, &1, composed), done && { done, fun }, :lists.reverse(accs), acc)
+      do_each(&Enumerable.reduce(enum, &1, composed), after_funs,
+              last && { last, fun }, :lists.reverse(accs), acc)
     end
 
-    defp do_each(_reduce, _done, _accs, { :halt, acc }) do
-      { :halted, acc }
-    end
-
-    defp do_each(reduce, done, accs, { :suspend, acc }) do
-      { :suspended, acc, &do_each(reduce, done, accs, &1) }
-    end
-
-    defp do_each(reduce, done, accs, { :cont, acc }) do
-      case reduce.({ :cont, [acc|accs] }) do
+    defp do_each(reduce, after_funs, last, accs, { command, acc }) do
+      case reduce.({ command, [acc|accs] }) do
         { :suspended, [acc|accs], continuation } ->
-          { :suspended, acc, &do_each(continuation, done, accs, &1) }
+          { :suspended, acc, &do_each(continuation, after_funs, last, accs, &1) }
         { :halted, [acc|_] } ->
+          lc fun inlist after_funs, do: fun.()
           { :halted, acc }
         { :done, [acc|_] = accs } ->
-          case done do
+          case last do
             nil ->
+              lc fun inlist after_funs, do: fun.()
               { :done, acc }
-            { done, fun } ->
-              case done.(fun).(accs) do
+            { last, fun } ->
+              res = case last.(fun).(accs) do
                 { :cont, [acc|_] }    -> { :done, acc }
                 { :halt, [acc|_] }    -> { :halted, acc }
                 { :suspend, [acc|_] } -> { :suspended, acc, &({ :done, &1 |> elem(1) }) }
               end
+              lc fun inlist after_funs, do: fun.()
+              res
           end
       end
     end
@@ -163,10 +160,30 @@ defmodule Stream do
   ## Transformers
 
   @doc """
-  Shortcut to `chunks(coll, n, n)`.
+  Executes the given function when the stream is done, halted
+  or an error happened during streaming. Useful for resource
+  clean up.
+
+  Callbacks registered later will be executed earlier.
+
+  ## Examples
+
+      iex> stream = Stream.after [1,2,3], fn -> Process.put(:done, true) end
+      iex> Enum.to_list(stream)
+      [1,2,3]
+      iex> Process.get(:done)
+      true
+
+  """
+  @spec unquote(:after)(Enumerable.t, (() -> term)) :: Enumerable.t
+  def unquote(:after)(Lazy[after: funs] = lazy, fun), do: lazy.after([fun|funs])
+  def unquote(:after)(enum, fun), do: Lazy[enum: enum, after: [fun]]
+
+  @doc """
+  Shortcut to `chunks(enum, n, n)`.
   """
   @spec chunks(Enumerable.t, non_neg_integer) :: Enumerable.t
-  def chunks(coll, n), do: chunks(coll, n, n, nil)
+  def chunks(enum, n), do: chunks(enum, n, n, nil)
 
   @doc """
   Streams the enumerable in chunks, containing `n` items each, where
@@ -487,6 +504,30 @@ defmodule Stream do
     lazy enum, fn(f1) -> R.reject(fun, f1) end
   end
 
+  @doc """
+  Runs the given stream.
+
+  This is useful when a stream needs to be run, for side effects,
+  and there is no interest in its return result.
+
+  ## Examples
+
+  Open up a file, replace all # by % and stream to another file
+  without loading the whole file in memory:
+
+      stream = File.stream!("code")
+      |> Stream.map(&String.replace(&1, "#", "%"))
+      |> File.stream_to!("new")
+
+  No computation will be done until we call one of the Enum functions
+  or `Stream.run/1`.
+  """
+  @spec run(Enumerable.t) :: :ok
+  def run(stream) do
+    Enumerable.reduce(stream, { :cont, nil }, fn(_, _) -> { :cont, nil } end)
+    :ok
+  end
+
   @doc """
   Lazily takes the next `n` items from the enumerable and stops
   enumeration.
@@ -911,8 +952,8 @@ defmodule Stream do
   defp lazy(enum, acc, fun),
     do: Lazy[enum: enum, funs: [fun], accs: [acc]]
 
-  defp lazy(Lazy[done: nil, funs: funs, accs: accs] = lazy, acc, fun, done),
-    do: lazy.funs([fun|funs]).accs([acc|accs]).done(done)
-  defp lazy(enum, acc, fun, done),
-    do: Lazy[enum: enum, funs: [fun], accs: [acc], done: done]
+  defp lazy(Lazy[last: nil, funs: funs, accs: accs] = lazy, acc, fun, last),
+    do: lazy.funs([fun|funs]).accs([acc|accs]).last(last)
+  defp lazy(enum, acc, fun, last),
+    do: Lazy[enum: enum, funs: [fun], accs: [acc], last: last]
 end

lib/elixir/test/elixir/file_test.exs

@@ -995,6 +995,23 @@ defmodule FileTest do
     end
   end
 
+  test :stream_to do
+    src  = fixture_path("file.txt")
+    dest = tmp_path("tmp_test.txt")
+
+    try do
+      stream = File.stream!(src)
+        |> Stream.map(&String.replace(&1, "O", "A"))
+        |> File.stream_to!(dest)
+
+      refute File.exists?(dest)
+      assert Stream.run(stream) == :ok
+      assert File.read(dest) == { :ok, "FAA\n" }
+    after
+      File.rm(dest)
+    end
+  end
+
   test :stream_bytes do
     src  = fixture_path("file.txt")
     dest = tmp_path("tmp_test.txt")

lib/elixir/test/elixir/stream_test.exs

@@ -25,6 +25,20 @@ defmodule StreamTest do
     assert Enum.to_list(stream) == [3,5,7]
   end
 
+  test "after" do
+    stream = Stream.after([1,2,3], fn -> Process.put(:after, true) end)
+
+    # Done
+    Process.put(:after, false)
+    assert Enum.to_list(stream) == [1,2,3]
+    assert Process.get(:after)
+
+    # Halted
+    Process.put(:after, false)
+    assert Enum.take(stream, 1) == [1]
+    assert Process.get(:after)
+  end
+
   test "chunks" do
     assert Stream.chunks([1, 2, 3, 4, 5], 2) |> Enum.to_list ==
            [[1, 2], [3, 4]]
@@ -395,6 +409,22 @@ defmodule StreamTest do
     assert Enum.zip(list, list) == Enum.zip(stream, stream)
   end
 
+  test "run" do
+    Process.put(:stream_each, [])
+    Process.put(:stream_after, false)
+
+    stream = [1,2,3]
+    |> Stream.after(fn -> Process.put(:stream_after, true) end)
+    |> Stream.each(fn x ->
+         Process.put(:stream_each, [x|Process.get(:stream_each)])
+       end)
+
+    assert is_lazy(stream)
+    assert Stream.run(stream) == :ok
+    assert Process.get(:stream_after)
+    assert Process.get(:stream_each) == [3,2,1]
+  end
+
   test "unfold only calculate values if needed" do
     stream = Stream.unfold(1, fn x -> if x > 0, do: {x, x-1}, else: throw(:boom) end)
     assert Enum.take(stream, 1) == [1]