3 files changed, 76 insertions, 8 deletions

diff --git a/printf.c b/printf.c
index 5c1add5..3e7df7d 100644
--- a/printf.c
+++ b/printf.c
@@ -69,6 +69,11 @@ static int bfs_printf_flush(CFILE *cfile, const struct bfs_printf *directive, co
 	return fflush(cfile->file);
 }
 
+/** Check if we can safely colorize this directive. */
+static bool should_color(CFILE *cfile, const struct bfs_printf *directive) {
+	return cfile->colors && strcmp(directive->str, "%s") == 0;
+}
+
 /**
  * Print a value to a temporary buffer before formatting it.
  */
@@ -210,7 +215,11 @@ static int bfs_printf_D(CFILE *cfile, const struct bfs_printf *directive, const
 
 /** %f: file name */
 static int bfs_printf_f(CFILE *cfile, const struct bfs_printf *directive, const struct BFTW *ftwbuf) {
-	return fprintf(cfile->file, directive->str, ftwbuf->path + ftwbuf->nameoff);
+	if (should_color(cfile, directive)) {
+		return cfprintf(cfile, "%pF", ftwbuf);
+	} else {
+		return fprintf(cfile->file, directive->str, ftwbuf->path + ftwbuf->nameoff);
+	}
 }
 
 /** %F: file system type */
@@ -273,14 +282,28 @@ static int bfs_printf_h(CFILE *cfile, const struct bfs_printf *directive, const
 		return -1;
 	}
 
-	int ret = fprintf(cfile->file, directive->str, buf);
+	int ret;
+	if (should_color(cfile, directive)) {
+		ret = cfprintf(cfile, "${di}%s${rs}", buf);
+	} else {
+		ret = fprintf(cfile->file, directive->str, buf);
+	}
+
 	free(copy);
 	return ret;
 }
 
 /** %H: current root */
 static int bfs_printf_H(CFILE *cfile, const struct bfs_printf *directive, const struct BFTW *ftwbuf) {
-	return fprintf(cfile->file, directive->str, ftwbuf->root);
+	if (should_color(cfile, directive)) {
+		if (ftwbuf->depth == 0) {
+			return cfprintf(cfile, "%pP", ftwbuf);
+		} else {
+			return cfprintf(cfile, "${di}%s${rs}", ftwbuf->root);
+		}
+	} else {
+		return fprintf(cfile->file, directive->str, ftwbuf->root);
+	}
 }
 
 /** %i: inode */
@@ -312,6 +335,10 @@ static int bfs_printf_l(CFILE *cfile, const struct bfs_printf *directive, const
 	const char *target = "";
 
 	if (ftwbuf->type == BFS_LNK) {
+		if (should_color(cfile, directive)) {
+			return cfprintf(cfile, "%pL", ftwbuf);
+		}
+
 		const struct bfs_stat *statbuf = bftw_cached_stat(ftwbuf, BFS_STAT_NOFOLLOW);
 		size_t len = statbuf ? statbuf->size : 0;
 
@@ -361,16 +388,32 @@ static int bfs_printf_n(CFILE *cfile, const struct bfs_printf *directive, const
 
 /** %p: full path */
 static int bfs_printf_p(CFILE *cfile, const struct bfs_printf *directive, const struct BFTW *ftwbuf) {
-	return fprintf(cfile->file, directive->str, ftwbuf->path);
+	if (should_color(cfile, directive)) {
+		return cfprintf(cfile, "%pP", ftwbuf);
+	} else {
+		return fprintf(cfile->file, directive->str, ftwbuf->path);
+	}
 }
 
 /** %P: path after root */
 static int bfs_printf_P(CFILE *cfile, const struct bfs_printf *directive, const struct BFTW *ftwbuf) {
-	const char *path = ftwbuf->path + strlen(ftwbuf->root);
-	if (path[0] == '/') {
-		++path;
+	size_t offset = strlen(ftwbuf->root);
+	if (ftwbuf->path[offset] == '/') {
+		++offset;
+	}
+
+	if (should_color(cfile, directive)) {
+		if (ftwbuf->depth == 0) {
+			return 0;
+		}
+
+		struct BFTW copybuf = *ftwbuf;
+		copybuf.path += offset;
+		copybuf.nameoff -= offset;
+		return cfprintf(cfile, "%pP", &copybuf);
+	} else {
+		return fprintf(cfile->file, directive->str, ftwbuf->path + offset);
 	}
-	return fprintf(cfile->file, directive->str, path);
 }
 
 /** %s: size */
diff --git a/tests.sh b/tests.sh
index df6eb04..7bcda74 100755
--- a/tests.sh
+++ b/tests.sh
@@ -765,6 +765,7 @@ bfs_tests=(
     test_printf_invalid_format
     test_printf_duplicate_flag
     test_printf_must_be_numeric
+    test_printf_color
 
     test_type_multi
 
@@ -2305,6 +2306,10 @@ function test_printf_must_be_numeric() {
     ! quiet invoke_bfs basic -printf '%+p'
 }
 
+function test_printf_color() {
+    LS_COLORS= bfs_diff -color -path './rainbow*' -printf '%H %h %f %p %P %l\n'
+}
+
 function test_fprintf() {
     invoke_bfs basic -fprintf scratch/test_fprintf.out '%%p(%p) %%d(%d) %%f(%f) %%h(%h) %%H(%H) %%P(%P) %%m(%m) %%M(%M) %%y(%y)\n'
     sort -o scratch/test_fprintf.out scratch/test_fprintf.out
diff --git a/tests/test_printf_color.out b/tests/test_printf_color.out
new file mode 100644
index 0000000..dc04d14
--- /dev/null
+++ b/tests/test_printf_color.out
@@ -0,0 +1,20 @@
+. . rainbow ./rainbow rainbow 
+. ./rainbow exec.sh ./rainbow/exec.sh rainbow/exec.sh 
+. ./rainbow socket ./rainbow/socket rainbow/socket 
+. ./rainbow broken ./rainbow/broken rainbow/broken nowhere
+. ./rainbow link.txt ./rainbow/link.txt rainbow/link.txt file.txt
+. ./rainbow sticky_ow ./rainbow/sticky_ow rainbow/sticky_ow 
+. ./rainbow sgid ./rainbow/sgid rainbow/sgid 
+. ./rainbow pipe ./rainbow/pipe rainbow/pipe 
+. ./rainbow ow ./rainbow/ow rainbow/ow 
+. ./rainbow sugid ./rainbow/sugid rainbow/sugid 
+. ./rainbow suid ./rainbow/suid rainbow/suid 
+. ./rainbow sticky ./rainbow/sticky rainbow/sticky 
+. ./rainbow file.dat ./rainbow/file.dat rainbow/file.dat 
+. ./rainbow file.txt ./rainbow/file.txt rainbow/file.txt 
+. ./rainbow mh1 ./rainbow/mh1 rainbow/mh1 
+. ./rainbow mh2 ./rainbow/mh2 rainbow/mh2 
+. ./rainbow star.gz ./rainbow/star.gz rainbow/star.gz 
+. ./rainbow star.tar ./rainbow/star.tar rainbow/star.tar 
+. ./rainbow star.tar.gz ./rainbow/star.tar.gz rainbow/star.tar.gz 
+. ./rainbow chardev_link ./rainbow/chardev_link rainbow/chardev_link /dev/null