From 932b482e4a0cf200741070784c4b0da8afd0fd22 Mon Sep 17 00:00:00 2001
From: Andrew Thompson <andrew@hijacked.us>
Date: Thu, 24 Jan 2013 12:56:26 -0500
Subject: [PATCH] Support printing iolists with ~s

---
 src/lager_format.erl   |  2 ++
 src/lager_trunc_io.erl | 10 ++++++++++
 test/trunc_io_eqc.erl  | 11 ++++++++---
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/lager_format.erl b/src/lager_format.erl
index 8084948..b834f27 100644
--- a/src/lager_format.erl
+++ b/src/lager_format.erl
@@ -29,6 +29,8 @@
 format(FmtStr, Args, MaxLen) ->
     format(FmtStr, Args, MaxLen, []).
 
+format([], [], _, _) ->
+    "";
 format(FmtStr, Args, MaxLen, Opts) when is_atom(FmtStr) ->
     format(atom_to_list(FmtStr), Args, MaxLen, Opts);
 format(FmtStr, Args, MaxLen, Opts) when is_binary(FmtStr) ->
diff --git a/src/lager_trunc_io.erl b/src/lager_trunc_io.erl
index c3a2d5b..0128c6a 100644
--- a/src/lager_trunc_io.erl
+++ b/src/lager_trunc_io.erl
@@ -394,6 +394,11 @@ alist([H|T], Max, Options) when H =:= $\t; H =:= $\n; H =:= $\r; H =:= $\v; H =:
 alist([H|T], Max, #print_options{force_strings=true} = Options) when is_integer(H) ->
     {L, Len} = alist(T, Max-1, Options),
     {[H|L], Len + 1};
+alist([H|T], Max, Options = #print_options{force_strings=true}) when is_binary(H) ->
+    {List, Len} = print(H, Max, Options),
+    %% no need to decrement depth, as we're in printable string mode
+    {Final, FLen} = alist(T, Max - Len, Options),
+    {[List|Final], FLen};
 alist(_, _, #print_options{force_strings=true}) ->
     erlang:error(badarg);
 alist([H|_L], _Max, _Options) ->
@@ -650,6 +655,11 @@ list_printing_test() ->
     ?assertEqual("[9|9]", lists:flatten(format("~p", [[9|9]], 50))),
     ok.
 
+iolist_printing_test() ->
+    ?assertEqual("iolist: HelloIamaniolist",
+        lists:flatten(format("iolist: ~s", [[$H, $e,  $l, $l, $o, "I", ["am", [<<"an">>], [$i, $o, $l, $i, $s, $t]]]], 1000))),
+    ok.
+
 tuple_printing_test() ->
     ?assertEqual("{}", lists:flatten(format("~p", [{}], 50))),
     ?assertEqual("{}", lists:flatten(format("~w", [{}], 50))),
diff --git a/test/trunc_io_eqc.erl b/test/trunc_io_eqc.erl
index 390b979..2c5c56b 100644
--- a/test/trunc_io_eqc.erl
+++ b/test/trunc_io_eqc.erl
@@ -66,7 +66,7 @@ gen_fmt_args() ->
                 "~~",
                 {"~10000000.p", gen_any(5)},
                 {"~w", gen_any(5)},
-                {"~s", oneof([gen_print_str(), gen_atom(), gen_quoted_atom(), gen_print_bin()])},
+                {"~s", oneof([gen_print_str(), gen_atom(), gen_quoted_atom(), gen_print_bin(), gen_iolist(5)])},
                 {"~1000000.P", gen_any(5), 4},
                 {"~W", gen_any(5), 4},
                 {"~i", gen_any(5)},
@@ -91,7 +91,7 @@ gen_fmt_args() ->
 
 %% Generates a printable string
 gen_print_str() ->
-    ?LET(Xs, list(char()), [X || X <- Xs, io_lib:printable_list([X]), X /= $~]).
+    ?LET(Xs, list(char()), [X || X <- Xs, io_lib:printable_list([X]), X /= $~, X < 255]).
 
 gen_print_bin() ->
     ?LET(Xs, gen_print_str(), list_to_binary(Xs)).
@@ -110,7 +110,12 @@ gen_any(MaxDepth) ->
            gen_fun()] ++
            [?LAZY(list(gen_any(MaxDepth - 1))) || MaxDepth /= 0] ++
            [?LAZY(gen_tuple(gen_any(MaxDepth - 1))) || MaxDepth /= 0]).
-          
+
+gen_iolist(0) ->
+    [];
+gen_iolist(Depth) ->
+    list(oneof([gen_char(), gen_print_str(), gen_print_bin(), gen_iolist(Depth-1)])).
+
 gen_atom() ->
     elements([abc, def, ghi]).