From 4094617f1ea15a389c484286be39836c9a478070 Mon Sep 17 00:00:00 2001
From: Wojtek Mach <wojtekmach@users.noreply.github.com>
Date: Thu, 20 Jun 2024 14:32:02 +0200
Subject: [PATCH] Accept module plugs in Kino.Proxy.listen/1 (#448)

---
 lib/kino/proxy.ex | 41 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 37 insertions(+), 4 deletions(-)

diff --git a/lib/kino/proxy.ex b/lib/kino/proxy.ex
index f9a84128..ec8e2419 100644
--- a/lib/kino/proxy.ex
+++ b/lib/kino/proxy.ex
@@ -72,14 +72,47 @@ defmodule Kino.Proxy do
   > ```
   """
 
+  @type plug() ::
+          (Plug.Conn.t() -> Plug.Conn.t())
+          | module()
+          | {module(), term()}
+
   @doc """
   Registers a request listener.
 
-  Expects the listener to be a function that handles a request
-  `Plug.Conn`.
+  Expects the listener to be a plug, that is, one of:
+
+    * a function plug: a `fun(conn)` function that takes a `Plug.Conn` and returns a `Plug.Conn`.
+
+    * a module plug: a `module` atom or a `{module, options}` tuple.
   """
-  @spec listen((Plug.Conn.t() -> Plug.Conn.t())) :: DynamicSupervisor.on_start_child()
-  def listen(fun) when is_function(fun, 1) do
+  @spec listen(plug()) :: DynamicSupervisor.on_start_child()
+  def listen(plug) do
+    fun =
+      case plug do
+        fun when is_function(fun, 1) ->
+          fun
+
+        mod when is_atom(mod) ->
+          opts = mod.init([])
+          &mod.call(&1, opts)
+
+        {mod, opts} when is_atom(mod) ->
+          opts = mod.init(opts)
+          &mod.call(&1, opts)
+
+        other ->
+          raise """
+          expected plug to be one of:
+
+            * fun(conn)
+            * module
+            * {module, options}
+
+          got: #{inspect(other)}
+          """
+      end
+
     case Kino.Bridge.get_proxy_handler_child_spec(fun) do
       {:ok, child_spec} ->
         Kino.start_child(child_spec)