Accept function in Phoenix.Socket.assign/2 (#6530)

* Add Phoenix.Socket.assign_lazy/3

* Refactoring

Author: timfjord

lib/phoenix/socket.ex

@@ -345,20 +345,29 @@ defmodule Phoenix.Socket do
 
   @doc """
   Adds key/value pairs to socket assigns.
+  Accepts a keyword list, a map, or a single-argument function.
 
-  A keyword list or a map of assigns must be given as argument to be merged into existing assigns.
+  When a keyword list or map is provided, it will be merged into the existing assigns.
+
+  If a function is given, it takes the current assigns as an argument and its return
+  value will be merged into the current assigns.
 
   ## Examples
 
       iex> assign(socket, name: "Elixir", logo: "💧")
       iex> assign(socket, %{name: "Elixir"})
+      iex> assign(socket, fn %{name: name, logo: logo} -> %{title: Enum.join([name, logo], " | ")} end)
 
   """
   def assign(%Socket{} = socket, keyword_or_map)
       when is_map(keyword_or_map) or is_list(keyword_or_map) do
     %{socket | assigns: Map.merge(socket.assigns, Map.new(keyword_or_map))}
   end
 
+  def assign(%Socket{} = socket, fun) when is_function(fun, 1) do
+    assign(socket, fun.(socket.assigns))
+  end
+
   @doc """
   Defines a channel matching the given topic and transports.
 

test/phoenix/socket/socket_test.exs

@@ -80,6 +80,15 @@ defmodule Phoenix.SocketTest do
       assert socket.assigns[:foo] == :baz
       assert socket.assigns[:abc] == :def
     end
+
+    test "accepts functions" do
+      socket = %Phoenix.Socket{}
+      assert socket.assigns[:foo] == nil
+      socket = assign(socket, :foo, :bar)
+      assert socket.assigns[:foo] == :bar
+      socket = assign(socket, fn %{foo: :bar} -> [baz: :quux] end)
+      assert socket.assigns[:baz] == :quux
+    end
   end
 
   describe "drainer_spec/1" do