Merge pull request #34 from barruumrex/master

Add headers option to encode to match decode

lib/csv/encoder.ex

@@ -16,6 +16,12 @@ defmodule CSV.Encoder do
 
     * `:separator`   – The separator token to use, defaults to `?,`. Must be a codepoint (syntax: ? + your separator token).
     * `:delimiter`   – The delimiter token to use, defaults to `\"\\r\\n\"`.
+    * `:headers`     – When set to `true`, uses the keys of the first map as the first element in the stream. All
+      subsequent elements are the values of the maps.
+      When set to a list, will use the given list as the first element in the stream and order all subsequent
+      elements using that list.
+      When set to `false` (default), will use the raw inputs as elements.
+      When set to anything but `false`, all elements in the input stream are assumed to be maps.
 
   ## Examples
 
@@ -26,25 +32,54 @@ defmodule CSV.Encoder do
       iex> Enum.take(2)
       [\"a,b\\r\\n\", \"c,d\\r\\n\"]
 
+  Convert a stream of maps into a stream of lines:
+
+      iex> [%{"a" => 1, "b" => 2}, %{"a" => 3, "b" => 4}] |>
+      iex> CSV.Encoder.encode(headers: true) |>
+      iex> Enum.to_list()
+      [\"a,b\\r\\n\", \"1,2\\r\\n\", \"3,4\\r\\n\"]
+
   Convert a stream of rows with cells with escape sequences into a stream of lines:
 
       iex> [[\"a\\nb\", \"\\tc\"], [\"de\", \"\\tf\\\"\"]] |>
-      iex> CSV.Encoder.encode(separator: ?\t, delimiter: \"\\n\") |>
+      iex> CSV.Encoder.encode(separator: ?\\t, delimiter: \"\\n\") |>
       iex> Enum.take(2)
-      [\"\\\"a\\nb\\\"\\t\\\"\\tc\\\"\\n\", \"de\\t\\\"\\tf\\\"\\\"\\\"\\n\"]
+      [\"\\\"a\\\\nb\\\"\\t\\\"\\\\tc\\\"\\n\", \"de\\t\\\"\\\\tf\\\"\\\"\\\"\\n\"]
   """
 
   def encode(stream, options \\ []) do
-    separator = options |> Keyword.get(:separator, @separator)
-    delimiter = options |> Keyword.get(:delimiter, @delimiter)
+    headers = options |> Keyword.get(:headers, false)
 
+    encode_stream(stream, headers, options)
+  end
+
+  defp encode_stream(stream, false, options) do
     stream |> Stream.transform(0, fn row, acc ->
-      {[ encode_row(row, separator, delimiter) <> delimiter ], acc + 1}
+      {[encode_row(row, options)], acc + 1} end)
+  end
+  defp encode_stream(stream, headers, options) do
+    stream |> Stream.transform(0, fn
+      row, 0 -> {[encode_row(get_headers(row, headers), options),
+                  encode_row(get_values(row, headers), options)], 1}
+      row, acc -> {[encode_row(get_values(row, headers), options)], acc + 1}
     end)
   end
 
-  defp encode_row(row, separator, delimiter) do
-    row |> Enum.map(&encode_cell(&1, separator, delimiter)) |> Enum.join(<< separator :: utf8 >>)
+  defp get_headers(row, true), do: Map.keys(row)
+  defp get_headers(_row, headers), do: headers
+
+  defp get_values(row, true), do: Map.values(row)
+  defp get_values(row, headers), do: headers |> Enum.map(&Map.get(row, &1))
+
+  defp encode_row(row, options) do
+    separator = options |> Keyword.get(:separator, @separator)
+    delimiter = options |> Keyword.get(:delimiter, @delimiter)
+
+    encoded = row
+      |> Enum.map(&encode_cell(&1, separator, delimiter))
+      |> Enum.join(<< separator :: utf8 >>)
+
+    encoded <> delimiter
   end
 
   defp encode_cell(cell, separator, delimiter) do

test/encoder_test.exs

@@ -1,6 +1,7 @@
 defmodule EncoderTest do
   use ExUnit.Case
   alias CSV.Encoder, as: Encoder
+  doctest Encoder
 
   test "encodes streams to csv strings" do
     result = Encoder.encode([~w(a b), ~w(c d)]) |> Enum.take(2)
@@ -27,4 +28,13 @@ defmodule EncoderTest do
     assert result == ["\"a\\t\"\t\"b\\re\"\n", "\"c\\tf\"\"\"\tdg\n"]
   end
 
+  test "use keys from first row as headers when headers: true" do
+    result = Encoder.encode([%{"a" => 1, "b" => 2}], headers: true) |> Enum.to_list()
+    assert result == ["a,b\r\n", "1,2\r\n"]
+  end
+
+  test "specified headers inserted as first row and used to order columns" do
+    result = Encoder.encode([%{"b" => 2}], headers: ["a", "b"]) |> Enum.to_list()
+    assert result == ["a,b\r\n", ",2\r\n"]
+  end
 end