Skip to content

Commit 2257758

Browse files
authored
common : change --color to accept on/off/auto, default to auto (#17827)
1 parent d9e03db commit 2257758

File tree

4 files changed

+54
-32
lines changed

4 files changed

+54
-32
lines changed

common/arg.cpp

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -708,6 +708,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
708708
params.use_jinja = true;
709709
}
710710

711+
params.use_color = tty_can_use_colors();
712+
711713
// load dynamic backends
712714
ggml_backend_load_all();
713715

@@ -790,10 +792,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
790792
}
791793
).set_examples({LLAMA_EXAMPLE_MAIN}));
792794
add_opt(common_arg(
793-
{"-co", "--color"},
794-
string_format("colorise output to distinguish prompt and user input from generations (default: %s)", params.use_color ? "true" : "false"),
795-
[](common_params & params) {
796-
params.use_color = true;
795+
{"-co", "--color"}, "[on|off|auto]",
796+
"Colorize output to distinguish prompt and user input from generations ('on', 'off', or 'auto', default: 'auto')\n"
797+
"'auto' enables colors when output is to a terminal",
798+
[](common_params & params, const std::string & value) {
799+
if (is_truthy(value)) {
800+
params.use_color = true;
801+
} else if (is_falsey(value)) {
802+
params.use_color = false;
803+
} else if (is_autoy(value)) {
804+
params.use_color = tty_can_use_colors();
805+
} else {
806+
throw std::invalid_argument(
807+
string_format("error: unknown value for --color: '%s'\n", value.c_str()));
808+
}
797809
}
798810
).set_examples({LLAMA_EXAMPLE_MAIN, LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP}));
799811
add_opt(common_arg(
@@ -1022,7 +1034,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
10221034
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO;
10231035
} else {
10241036
throw std::runtime_error(
1025-
string_format("error: unkown value for --flash-attn: '%s'\n", value.c_str()));
1037+
string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str()));
10261038
}
10271039
}).set_env("LLAMA_ARG_FLASH_ATTN"));
10281040
add_opt(common_arg(
@@ -2696,7 +2708,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
26962708
common_log_set_colors(common_log_main(), LOG_COLORS_AUTO);
26972709
} else {
26982710
throw std::invalid_argument(
2699-
string_format("error: unkown value for --log-colors: '%s'\n", value.c_str()));
2711+
string_format("error: unknown value for --log-colors: '%s'\n", value.c_str()));
27002712
}
27012713
}
27022714
).set_env("LLAMA_LOG_COLORS"));

common/common.cpp

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -982,6 +982,32 @@ std::vector<common_file_info> fs_list(const std::string & path, bool include_dir
982982
return files;
983983
}
984984

985+
//
986+
// TTY utils
987+
//
988+
989+
bool tty_can_use_colors() {
990+
// Check NO_COLOR environment variable (https://no-color.org/)
991+
if (const char * no_color = std::getenv("NO_COLOR")) {
992+
if (no_color[0] != '\0') {
993+
return false;
994+
}
995+
}
996+
997+
// Check TERM environment variable
998+
if (const char * term = std::getenv("TERM")) {
999+
if (std::strcmp(term, "dumb") == 0) {
1000+
return false;
1001+
}
1002+
}
1003+
1004+
// Check if stdout and stderr are connected to a terminal
1005+
// We check both because log messages can go to either
1006+
bool stdout_is_tty = isatty(fileno(stdout));
1007+
bool stderr_is_tty = isatty(fileno(stderr));
1008+
1009+
return stdout_is_tty || stderr_is_tty;
1010+
}
9851011

9861012
//
9871013
// Model utils

common/common.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -655,6 +655,13 @@ struct common_file_info {
655655
};
656656
std::vector<common_file_info> fs_list(const std::string & path, bool include_directories);
657657

658+
//
659+
// TTY utils
660+
//
661+
662+
// Auto-detect if colors can be enabled based on terminal and environment
663+
bool tty_can_use_colors();
664+
658665
//
659666
// Model utils
660667
//

common/log.cpp

Lines changed: 3 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
#include "common.h"
12
#include "log.h"
23

34
#include <chrono>
@@ -26,30 +27,6 @@ void common_log_set_verbosity_thold(int verbosity) {
2627
common_log_verbosity_thold = verbosity;
2728
}
2829

29-
// Auto-detect if colors should be enabled based on terminal and environment
30-
static bool common_log_should_use_colors_auto() {
31-
// Check NO_COLOR environment variable (https://no-color.org/)
32-
if (const char * no_color = std::getenv("NO_COLOR")) {
33-
if (no_color[0] != '\0') {
34-
return false;
35-
}
36-
}
37-
38-
// Check TERM environment variable
39-
if (const char * term = std::getenv("TERM")) {
40-
if (std::strcmp(term, "dumb") == 0) {
41-
return false;
42-
}
43-
}
44-
45-
// Check if stdout and stderr are connected to a terminal
46-
// We check both because log messages can go to either
47-
bool stdout_is_tty = isatty(fileno(stdout));
48-
bool stderr_is_tty = isatty(fileno(stderr));
49-
50-
return stdout_is_tty || stderr_is_tty;
51-
}
52-
5330
static int64_t t_us() {
5431
return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
5532
}
@@ -391,7 +368,7 @@ struct common_log * common_log_main() {
391368
static std::once_flag init_flag;
392369
std::call_once(init_flag, [&]() {
393370
// Set default to auto-detect colors
394-
log.set_colors(common_log_should_use_colors_auto());
371+
log.set_colors(tty_can_use_colors());
395372
});
396373

397374
return &log;
@@ -422,7 +399,7 @@ void common_log_set_file(struct common_log * log, const char * file) {
422399

423400
void common_log_set_colors(struct common_log * log, log_colors colors) {
424401
if (colors == LOG_COLORS_AUTO) {
425-
log->set_colors(common_log_should_use_colors_auto());
402+
log->set_colors(tty_can_use_colors());
426403
return;
427404
}
428405

0 commit comments

Comments
 (0)