livebook/lib/livebook_cli.ex

defmodule LivebookCLI do
  @switches [
    help: :boolean,
    version: :boolean
  ]

  @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

    case OptionParser.parse(args, strict: @switches, aliases: @aliases) do
      {[help: true], [], _} ->
        display_help()

      {[version: true], _, _} ->
        display_version()

      {[help: true], [name], _} ->
        LivebookCLI.Command.usage(name)

      {_, [name | _], _} ->
        # We want to keep the switches for the command, which is being ignored here
        LivebookCLI.Command.call(name, List.delete(args, name))
    end
  end

  defp unix?(), do: match?({:unix, _}, :os.type())

  defp display_help() do
    info("""
    Livebook is an interactive notebook system for Elixir

    Usage: livebook [command] [options]

    Available commands:

      livebook server    Starts the Livebook web application
      livebook deploy    Deploys a notebook to Livebook Teams

    The --help and --version options can be given instead of a command for usage and versioning information.\
    """)
  end

  defp display_version() do
    info("""
    #{:erlang.system_info(:system_version)}
    Elixir #{System.build_info()[:build]}
    Livebook #{Livebook.Config.app_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"

    # 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

  @doc """
  Logs an error message.
  """
  def error(message) do
    message
    |> put_level(:error)
    |> IO.ANSI.format()
    |> IO.puts()
  end

  @doc """
  Logs an info message.
  """
  def info(message) do
    message
    |> put_level(:info)
    |> IO.ANSI.format()
    |> IO.puts()
  end

  @doc """
  Logs a debug message.
  """
  if Mix.env() == :dev do
    def debug(message) do
      message
      |> put_level(:debug)
      |> IO.ANSI.format()
      |> IO.puts()
    end
  else
    def debug(message) do
      _ = put_level(message, :debug)
      :ok
    end
  end

  @doc """
  Logs a warning message.
  """
  def warning(message) do
    message
    |> put_level(:warning)
    |> IO.ANSI.format()
    |> IO.warn()
  end

  @doc """
  Logs an error message.
  """
  def raise(exception, command_name, stacktrace) when is_exception(exception) do
    exception
    |> format_exception(command_name, stacktrace)
    |> IO.ANSI.format()
    |> IO.puts()

    System.halt(1)
  end

  defp put_level(message, :info) do
    ["[info] ", message]
  end

  defp put_level(message, :error) do
    [:red, "[error] ", message]
  end

  defp put_level(message, :warning) do
    ["[warning] ", message]
  end

  defp put_level(message, :debug) do
    [:cyan, "[debug] ", message]
  end

  defp format_exception(%OptionParser.ParseError{} = exception, command_name, _) do
    [
      :red,
      :bright,
      """
      #{Exception.message(exception)}

      For more information try:

         livebook #{command_name} --help
      """
    ]
  end

  defp format_exception(exception, _command_name, stacktrace) do
    [
      :red,
      :bright,
      Exception.format(:error, exception, stacktrace)
    ]
  end
end