Improve CLI entrypoint

Author: aleDsz

lib/livebook_cli.ex

@@ -1,129 +1,47 @@
 defmodule LivebookCLI do
-  def usage() do
-    """
-    Usage: livebook [command] [options]
+  alias LivebookCLI.{Task, Utils}
 
-    Available commands:
+  @switches [
+    help: :boolean,
+    version: :boolean
+  ]
 
-      livebook server    Starts the Livebook web application
-
-    The --help and --version options can be given instead of a command for usage and versioning information.
-    """
-  end
+  @aliases [
+    h: :help,
+    v: :version
+  ]
 
   def main(args) do
-    {:ok, _} = Application.ensure_all_started(:elixir)
-
-    extract_priv!()
-
-    :ok = Application.load(:livebook)
-
-    if unix?() do
-      Application.put_env(:elixir, :ansi_enabled, true)
-    end
-
-    call(args)
-  end
-
-  defp unix?(), do: match?({:unix, _}, :os.type())
-
-  defp call([arg]) when arg in ["--help", "-h"], do: display_help()
-  defp call([arg]) when arg in ["--version", "-v"], do: display_version()
-
-  defp call([task_name | args]) do
-    case find_task(task_name) do
-      nil ->
-        IO.ANSI.format([:red, "Unknown command #{task_name}\n"]) |> IO.puts()
-        IO.write(usage())
-
-      task ->
-        call_task(task, args)
-    end
-  end
-
-  defp call(_args), do: IO.write(usage())
-
-  defp find_task("server"), do: LivebookCLI.Server
-  defp find_task(_), do: nil
-
-  defp call_task(task, [arg]) when arg in ["--help", "-h"] do
-    IO.write(task.usage())
-  end
-
-  defp call_task(task, args) do
-    try do
-      task.call(args)
-    rescue
-      error in OptionParser.ParseError ->
-        IO.ANSI.format([
-          :red,
-          Exception.message(error),
-          "\n\nFor more information try --help"
-        ])
-        |> IO.puts()
-
-      error ->
-        IO.ANSI.format([:red, Exception.format(:error, error, __STACKTRACE__), "\n"]) |> IO.puts()
+    Utils.setup()
+
+    case Utils.option_parse(args, strict: @switches, aliases: @aliases) do
+      {parsed, [], _} when parsed.help -> display_help()
+      {parsed, [name], _} when parsed.help -> Task.usage(name)
+      {parsed, _, _} when parsed.version -> display_version()
+      # We want to keep the switches for the task
+      {_, [name | _], _} -> Task.call(name, List.delete(args, name))
     end
   end
 
   defp display_help() do
-    IO.puts("Livebook is an interactive notebook system for Elixir\n")
-    IO.write(usage())
-  end
-
-  defp display_version() do
-    IO.puts(:erlang.system_info(:system_version))
-    IO.puts("Elixir " <> System.build_info()[:build])
-
-    version = Livebook.Config.app_version()
-    IO.puts("\nLivebook #{version}")
-  end
-
-  import Record
-  defrecord(:zip_file, extract(:zip_file, from_lib: "stdlib/include/zip.hrl"))
-
-  defp extract_priv!() do
-    archive_dir = Path.join(Livebook.Config.tmp_path(), "escript")
-    extracted_path = Path.join(archive_dir, "extracted")
-    in_archive_priv_path = ~c"livebook/priv"
+    Utils.print_text("""
+    Livebook is an interactive notebook system for Elixir
 
-    # In dev we want to extract fresh directory on every boot
-    if Livebook.Config.app_version() =~ "-dev" do
-      File.rm_rf!(archive_dir)
-    end
-
-    # When temporary directory is cleaned by the OS, the directories
-    # may be left in place, so we use a regular file (extracted) to
-    # check if the extracted archive is already available
-    if not File.exists?(extracted_path) do
-      {:ok, sections} = :escript.extract(:escript.script_name(), [])
-      archive = Keyword.fetch!(sections, :archive)
-
-      file_filter = fn zip_file(name: name) ->
-        List.starts_with?(name, in_archive_priv_path)
-      end
-
-      case :zip.extract(archive, cwd: String.to_charlist(archive_dir), file_filter: file_filter) do
-        {:ok, _} ->
-          :ok
+    Usage: livebook [command] [options]
 
-        {:error, error} ->
-          print_error_and_exit(
-            "Livebook failed to extract archive files, reason: #{inspect(error)}"
-          )
-      end
+    Available commands:
 
-      File.touch!(extracted_path)
-    end
+      livebook server    Starts the Livebook web application
 
-    priv_dir = Path.join(archive_dir, in_archive_priv_path)
-    Application.put_env(:livebook, :priv_dir, priv_dir, persistent: true)
+    The --help and --version options can be given instead of a command for usage and versioning information.\
+    """)
   end
 
-  @spec print_error_and_exit(String.t()) :: no_return()
-  defp print_error_and_exit(message) do
-    IO.ANSI.format([:red, message]) |> IO.puts()
-    System.halt(1)
+  defp display_version() do
+    Utils.print_text("""
+    #{:erlang.system_info(:system_version)}
+    Elixir #{System.build_info()[:build]}
+    Livebook #{Livebook.Config.app_version()}\
+    """)
   end
 end

lib/livebook_cli/task.ex

@@ -1,11 +1,44 @@
 defmodule LivebookCLI.Task do
+  import LivebookCLI.Utils
+
   @doc """
   Returns a description of the task usage.
   """
-  @callback usage() :: String.t()
+  @callback usage() :: IO.chardata()
 
   @doc """
   Runs the task with the given list of command line arguments.
   """
   @callback call(args :: list(String.t())) :: :ok
+
+  @doc """
+  Runs the task with the given list of command line arguments.
+  """
+  @spec call(String.t(), list(String.t())) :: :ok
+  def call(name, args) do
+    task = fetch_task!(name)
+    task.call(args)
+
+    :ok
+  rescue
+    exception -> log_exception(exception, name, __STACKTRACE__)
+  end
+
+  @doc """
+  Shows the description of the task usage.
+  """
+  @spec usage(String.t()) :: :ok
+  def usage(name) do
+    task = fetch_task!(name)
+    print_text(task.usage())
+
+    :ok
+  rescue
+    exception -> log_exception(exception, name, __STACKTRACE__)
+  end
+
+  @spec fetch_task!(String.t()) :: module() | no_return()
+  defp fetch_task!("server"), do: LivebookCLI.Server
+  defp fetch_task!("deploy"), do: LivebookCLI.Deploy
+  defp fetch_task!(name), do: raise("Unknown command #{name}")
 end

lib/livebook_cli/utils.ex

@@ -0,0 +1,112 @@
+defmodule LivebookCLI.Utils do
+  def setup do
+    {:ok, _} = Application.ensure_all_started(:elixir)
+
+    extract_priv!()
+
+    :ok = Application.load(:livebook)
+
+    if unix?() do
+      Application.put_env(:elixir, :ansi_enabled, true)
+    end
+  end
+
+  defp unix?(), do: match?({:unix, _}, :os.type())
+
+  def option_parse(argv, opts \\ []) do
+    {parsed, argv, errors} = OptionParser.parse(argv, opts)
+    {Enum.into(parsed, %{}), argv, errors}
+  end
+
+  def log_info(message) do
+    IO.puts(message)
+  end
+
+  if Mix.env() == :dev do
+    def log_debug(message) do
+      [:cyan, message]
+      |> IO.ANSI.format()
+      |> IO.puts()
+    end
+  else
+    def log_debug(_message) do
+      :ok
+    end
+  end
+
+  def log_warning(message) do
+    [:yellow, message]
+    |> IO.ANSI.format()
+    |> IO.warn()
+  end
+
+  def print_text(message) do
+    message
+    |> IO.ANSI.format()
+    |> IO.puts()
+  end
+
+  @spec log_exception(Exception.t(), String.t(), Exception.stacktrace()) :: no_return()
+  def log_exception(exception, command_name, stacktrace) when is_exception(exception) do
+    [:red, format_exception(exception, command_name, stacktrace)]
+    |> IO.ANSI.format()
+    |> IO.puts()
+
+    System.halt(1)
+  end
+
+  defp format_exception(%OptionParser.ParseError{} = exception, command_name, _) do
+    """
+    #{Exception.message(exception)}
+
+    For more information try:
+
+       livebook #{command_name} --help
+    """
+  end
+
+  defp format_exception(%RuntimeError{} = exception, _, _) do
+    Exception.message(exception)
+  end
+
+  defp format_exception(exception, _, stacktrace) do
+    Exception.format(:error, exception, stacktrace)
+  end
+
+  import Record
+  defrecord(:zip_file, extract(:zip_file, from_lib: "stdlib/include/zip.hrl"))
+
+  defp extract_priv!() do
+    archive_dir = Path.join(Livebook.Config.tmp_path(), "escript")
+    extracted_path = Path.join(archive_dir, "extracted")
+    in_archive_priv_path = ~c"livebook/priv"
+
+    # In dev we want to extract fresh directory on every boot
+    if Livebook.Config.app_version() =~ "-dev" do
+      File.rm_rf!(archive_dir)
+    end
+
+    # When temporary directory is cleaned by the OS, the directories
+    # may be left in place, so we use a regular file (extracted) to
+    # check if the extracted archive is already available
+    if not File.exists?(extracted_path) do
+      {:ok, sections} = :escript.extract(:escript.script_name(), [])
+      archive = Keyword.fetch!(sections, :archive)
+
+      file_filter = fn zip_file(name: name) ->
+        List.starts_with?(name, in_archive_priv_path)
+      end
+
+      opts = [cwd: String.to_charlist(archive_dir), file_filter: file_filter]
+
+      with {:error, error} <- :zip.extract(archive, opts) do
+        raise "Livebook failed to extract archive files, reason: #{inspect(error)}"
+      end
+
+      File.touch!(extracted_path)
+    end
+
+    priv_dir = Path.join(archive_dir, in_archive_priv_path)
+    Application.put_env(:livebook, :priv_dir, priv_dir, persistent: true)
+  end
+end