Create ansi.lua and implement --no-ansi

Author: gitpushjoe

core/ansi.lua

@@ -0,0 +1,54 @@
+local M = {}
+
+---@param code number
+---@param close_code number?
+local apply_code = function(code, close_code)
+	close_code = close_code or 39
+	---@param str string
+	---@return string
+	return function(str)
+		return string.format("\27[%dm%s\27[%dm", code, str, close_code)
+	end
+end
+
+---@alias ANSI_Func (fun(s: string): string)|({ __call: ANSI_Func})
+
+---@param code number
+---@return { __call: ANSI_Func; bg: ANSI_Func; bright: ANSI_Func; bright_bg: ANSI_Func}
+local function color(code)
+	return setmetatable({
+		bg = apply_code(code + 10, 49),
+		bright = apply_code(code + 60),
+		bright_bg = apply_code(code + 70, 49),
+	}, {
+		__call = function(_, str)
+			return apply_code(code)(str)
+		end,
+	})
+end
+
+---@param func1 ANSI_Func
+---@param func2 ANSI_Func
+---@return ANSI_Func
+M.compose = function(func1, func2)
+	---@param str string
+	---@return string
+	return function(str)
+		return func1(func2(str))
+	end
+end
+
+M.black = color(30)
+M.red = color(31)
+M.green = color(32)
+M.yellow = color(33)
+M.blue = color(34)
+M.magenta = color(35)
+M.cyan = color(36)
+M.white = color(37)
+
+M.none = apply_code(0, 0)
+M.bold = apply_code(1, 22)
+M.italic = apply_code(3, 23)
+
+return M

core/context.lua

@@ -54,6 +54,9 @@ local Config = require("core.config")
 --- associated with the mock filesystem. Otherwise, will be the Lua `io`
 --- library.
 --- @field io iolib
+---
+--- Specifies if `--no-ansi` was not passed.
+--- @field ansi_enabled boolean
 
 local Context = {}
 Context.__index = Context
@@ -86,6 +89,7 @@ Context.errors = {
 ---@param plan_stream Stream
 ---@param text_stream Stream
 ---@param preserve boolean
+---@param ansi_enabled boolean
 ---@return Context?, string?
 function Context:new(
 	config,
@@ -97,7 +101,8 @@ function Context:new(
 	auto_confirm,
 	plan_stream,
 	text_stream,
-	preserve
+	preserve,
+	ansi_enabled
 )
 	self = {}
 	---@cast self Context
@@ -125,6 +130,7 @@ function Context:new(
 		{ auto_confirm, "boolean", "auto_confirm" },
 		{ preserve, "boolean", "preserve" },
 		{ is_dry_run, "boolean", "is_dry_run" },
+		{ ansi_enabled, "boolean", "ansi_enabled" },
 	})
 	if err then
 		return nil, err
@@ -148,6 +154,7 @@ function Context:new(
 	self.lines = {}
 	self.io = mock_filesystem and mock_filesystem.io or io
 	self.preserve = preserve
+	self.ansi_enabled = ansi_enabled
 	self.is_dry_run = is_dry_run
 
 	local lines = {}

core/fold.lua

@@ -260,7 +260,7 @@ end
 M.execute = function(section_root, ctx, is_dry_run)
 	is_dry_run = is_dry_run or false
 	local io = ctx.io
-	local plan = Plan:new()
+	local plan = Plan:new(ctx.ansi_enabled)
 	local created_or_modified_paths = {}
 
 	---@param path_str string?

core/plan/action.lua

@@ -1,11 +1,12 @@
 local utils = require("core.utils")
 local Path = require("core.path")
+local ansi = require("core.ansi")
 
----@alias CREATE_ACTION { type: "CREATE", path: Path, lines: string[], tostring: fun(self: Action): string }
----@alias OVERWRITE_ACTION { type: "OVERWRITE", path: Path, lines: string[], tostring: fun(self: Action): string }
----@alias MKDIR_ACTION { type: "MKDIR", path: Path, tostring: fun(self: Action): string }
----@alias IGNORE_ACTION { type: "IGNORE", path: Path, lines: string[], tostring: fun(self: Action): string }
----@alias RENAME_ACTION { type: "RENAME", path: Path, new_path: Path }
+---@alias CREATE_ACTION { type: "CREATE", path: Path, lines: string[], tostring: fun(self: Action, enable_ansi: boolean?): string }
+---@alias OVERWRITE_ACTION { type: "OVERWRITE", path: Path, lines: string[], tostring: fun(self: Action, enable_ansi: boolean?): string }
+---@alias MKDIR_ACTION { type: "MKDIR", path: Path, tostring: fun(self: Action, enable_ansi: boolean?): string }
+---@alias IGNORE_ACTION { type: "IGNORE", path: Path, lines: string[], tostring: fun(self: Action, enable_ansi: boolean?): string }
+---@alias RENAME_ACTION { type: "RENAME", path: Path, new_path: Path, tostring: fun(self: Action, enable_ansi: boolean?): string }
 
 ---@alias Action CREATE_ACTION|OVERWRITE_ACTION|MKDIR_ACTION|IGNORE_ACTION|RENAME_ACTION
 
@@ -17,21 +18,26 @@ M.MKDIR = "MKDIR"
 M.IGNORE = "IGNORE"
 M.RENAME = "RENAME"
 
---- TODO(gitpushjoe): add option for no colors
 --- @param self Action
+--- @param enable_ansi boolean
 --- @return string
-local tostring = function(self)
-	local text = self.type == M.CREATE and "\27[32m"
-		or self.type == M.MKDIR and "\27[33m"
-		or self.type == M.OVERWRITE and "\27[35m"
-		or self.type == M.IGNORE and "\27[31m"
-		or self.type == M.RENAME and "\27[96m"
-		or error("Unknown action type: " .. self.type)
-	text = text
-		.. "[ "
-		.. string.rep(" ", #"OVERWRITE" - #self.type)
-		.. self.type
-		.. " ]  "
+local tostring = function(self, enable_ansi)
+	local color = enable_ansi
+			and (({
+				[M.CREATE] = ansi.green,
+				[M.MKDIR] = ansi.yellow,
+				[M.OVERWRITE] = ansi.magenta,
+				[M.IGNORE] = ansi.red,
+				[M.RENAME] = ansi.cyan,
+			})[self.type] or error("Unknown action type: " .. self.type))
+		or ansi.none
+	local bold = enable_ansi and ansi.bold or ansi.none
+	local text = bold(
+		"[ "
+			.. string.rep(" ", #"OVERWRITE" - #self.type)
+			.. self.type
+			.. " ]  "
+	)
 	local line_count_str = "-"
 	local char_count_str = "-"
 	if self.lines then
@@ -42,19 +48,21 @@ local tostring = function(self)
 		end
 		char_count_str = tostring(char_count)
 	end
-	return text
-		.. string.rep(" ", #"lines " - #line_count_str)
-		.. line_count_str
-		.. "  "
-		.. string.rep(" ", #"chars " - #char_count_str)
-		.. char_count_str
-		.. "  "
-		.. tostring(self.path)
-		.. (self.type == M.RENAME and "\n" .. string.rep(
-			" ",
-			#"action         lines   chars   " - #"->  "
-		) .. "->  " .. tostring(self.new_path) or "")
-		.. "\27[0m"
+	return color(
+		text
+			.. string.rep(" ", #"lines " - #line_count_str)
+			.. line_count_str
+			.. "  "
+			.. string.rep(" ", #"chars " - #char_count_str)
+			.. char_count_str
+			.. "  "
+			.. tostring(self.path)
+			.. (self.type == M.RENAME and "\n" .. string.rep(
+				" ",
+				#"action         lines   chars   " - #"->  "
+			) .. "->  " .. tostring(self.new_path) or "")
+			.. "\27[0m"
+	)
 end
 
 ---@param path Path

core/plan/plan.lua

@@ -1,14 +1,19 @@
 ---@class Plan
 ---@field actions Action[]
+---@field enable_ansi boolean
 Plan = {}
 Plan.__index = Plan
 Plan.__name = "Plan"
 
+---@param enable_ansi boolean?
 ---@return Plan
-function Plan:new()
+function Plan:new(enable_ansi)
+	enable_ansi = enable_ansi == nil and true or enable_ansi
 	self = {}
 	setmetatable(self, Plan)
+	---@cast self Plan
 	self.actions = {}
+	self.enable_ansi = not not enable_ansi
 	return self
 end
 
@@ -23,7 +28,7 @@ end
 function Plan:__tostring()
 	local text = "action         lines   chars   path"
 	for _, action in ipairs(self.actions) do
-		text = text .. "\n" .. action:tostring()
+		text = text .. "\n" .. action:tostring(self.enable_ansi)
 	end
 	return text
 end

core/section.lua

@@ -231,9 +231,9 @@ function Section:__tostring()
 		.. "}"
 		.. ",\n\tpath = "
 		.. (self.path and ('"' .. tostring(self.path) .. '"') or "nil")
-		.. ",\n\tlines = {"
+		.. ",\n\tlines = "
 		.. (function()
-			local text = ""
+			local text = "{"
 			local lines = self:get_lines()
 			for i, line in ipairs(lines) do
 				if line then
@@ -248,7 +248,7 @@ function Section:__tostring()
 						.. (i ~= #lines and ", " or "")
 				end
 			end
-			return text
+			return text .. "}"
 		end)()
 		.. ",\n\ttransformed_lines = "
 		.. (function()
@@ -260,18 +260,19 @@ function Section:__tostring()
 			for i, line in ipairs(lines) do
 				if line then
 					text = text
-					.. '"'
-					.. line:gsub("\\", "\\\\")
-						:gsub("\n", "\\n")
-						:gsub("\r", "\\r")
-						:gsub("\t", "\\t")
-					:gsub('"', '\\"')
-					.. '"'
-					.. (i ~= #lines and ", " or "")
+						.. '"'
+						.. line:gsub("\\", "\\\\")
+							:gsub("\n", "\\n")
+							:gsub("\r", "\\r")
+							:gsub("\t", "\\t")
+							:gsub('"', '\\"')
+						.. '"'
+						.. (i ~= #lines and ", " or "")
 				end
 			end
 			return text .. "}"
 		end)()
+		.. "\n}"
 end
 
 return Section

main.lua

@@ -20,14 +20,7 @@ local custom_configs = require("configs")
 local Parser = require("core.parser")
 local streams = require("core.streams")
 local utils = require("core.utils")
-
----@param err string?
-local error = function(err)
-	io.stderr:write(
-		"\27[1;31m" .. "ERROR: \27[0;91m" .. (err or "") .. "\n \27[0m"
-	)
-	os.exit(1)
-end
+local ansi = require("core.ansi")
 
 ---@param text string
 ---@param stream Stream
@@ -61,6 +54,11 @@ local parser = assert(Parser:new({
 		"-p",
 		"Will not edit the source file.",
 	},
+	{
+		"--no-ansi",
+		"-na",
+		"Disables ANSI coloring."
+	},
 }, {
 	{
 		"--plan-stream <stream>",
@@ -84,11 +82,23 @@ local parser = assert(Parser:new({
 	},
 }))
 
+local ansi_enabled = false
+
+---@param err string?
+local error = function(err)
+	local red = ansi_enabled and ansi.red or ansi.none
+	local bold = ansi_enabled and ansi.bold or ansi.none
+	io.stderr:write(red("ERROR: " .. bold(err or "")))
+	os.exit(1)
+end
+
 local success, err = parser:parse(arg)
 if not success then
 	error(err)
 end
 
+ansi_enabled = parser:find("--no-ansi") == nil
+
 if parser:find("--help") then
 	print(parser:get_helptext())
 	os.exit(0)
@@ -142,7 +152,8 @@ ctx, err = Context:new(
 	parser:find("--yes") ~= nil,
 	plan_stream or streams.STDOUT,
 	text_stream or streams.STDOUT,
-	parser:find("--preserve") ~= nil
+	parser:find("--preserve") ~= nil,
+	ansi_enabled
 )
 
 if not ctx then

readme.md

@@ -186,6 +186,7 @@ To run crazywall on a file, you can simply run `$ cw <file>`. You can add any of
 |`--config <config>`, `-c <config>`|Uses the config named <config> in `configs.lua`. Defaults to `"DEFAULT"`.
 |`--dry-run`, `-dr`|Enable dry-run, which will not modify or add any new files or directories.
 |`--help`, `-h`|Prints the helptext.
+|`--no-ansi`, `-na`|Disables ANSI coloring.
 |`--out <file>`, `-o <file>`|Sets the destination for the new source text to `<file>`. Defaults to the path to the source file.
 |`--plan_stream <stream>`, `-ps <stream>`|The stream to print the crazywall plan object to. (0 for none, 1 for stdout, 2 for stderr.)  Defaults to 1.
 |`--preserve`, `-p`|Will not edit the source file.