Accept function in Phoenix.Component.assign/2 (#4051)

Author: timfjord

lib/phoenix_component.ex

@@ -1433,12 +1433,15 @@ defmodule Phoenix.Component do
 
   The first argument is either a LiveView `socket` or an `assigns` map from function components.
 
-  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 as the second argument, 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_or_assigns, keyword_or_map)
@@ -1448,6 +1451,16 @@ defmodule Phoenix.Component do
     end)
   end
 
+  def assign(socket_or_assigns, fun) when is_function(fun, 1) do
+    assigns =
+      case socket_or_assigns do
+        %Socket{assigns: assigns} -> assigns
+        assigns -> assigns
+      end
+
+    assign(socket_or_assigns, fun.(assigns))
+  end
+
   defp validate_assign_key!(:flash) do
     raise ArgumentError,
           ":flash is a reserved assign by LiveView and it cannot be set directly. " <>

test/phoenix_component_test.exs

@@ -67,6 +67,15 @@ defmodule Phoenix.ComponentUnitTest do
       assert socket.assigns.existing == [:foo, :bam]
       assert socket.assigns.__changed__.existing == [:foo, :bar]
     end
+
+    test "allows functions" do
+      socket = assign(@socket, fn _ -> [existing: [:foo, :bar]] end)
+      socket = Utils.clear_changed(socket)
+
+      socket = assign(socket, fn %{existing: [:foo, :bar]} -> %{existing: [:foo, :baz]} end)
+      assert socket.assigns.existing == [:foo, :baz]
+      assert socket.assigns.__changed__.existing == [:foo, :bar]
+    end
   end
 
   describe "assign with assigns" do
@@ -104,6 +113,16 @@ defmodule Phoenix.ComponentUnitTest do
       assert assigns.map == %{foo: :baz}
       assert assigns.__changed__ == nil
     end
+
+    test "allows functions" do
+      assigns = assign(@assigns_changes, fn _ -> %{map: %{foo: :baz}} end)
+      assert assigns.map == %{foo: :baz}
+      assert assigns.__changed__[:map] == %{foo: :bar}
+
+      assigns = assign(@assigns_nil_changes, fn _ -> [map: %{foo: :baz}] end)
+      assert assigns.map == %{foo: :baz}
+      assert assigns.__changed__ == nil
+    end
   end
 
   describe "assign_new with socket" do