diff --git a/src/pp/containers_pp.ml b/src/pp/containers_pp.ml
index 57f73135..a5f06df2 100644
--- a/src/pp/containers_pp.ml
+++ b/src/pp/containers_pp.ml
@@ -1,3 +1,33 @@
+module B = Buffer
+
+module Ext = struct
+ module type OUT = sig
+ val char : char -> unit
+ val sub_string : string -> int -> int -> unit
+ val string : string -> unit
+ val newline : unit -> unit
+ end
+
+ type out = (module OUT)
+
+ let out_of_buf (buf : Buffer.t) : out =
+ (module struct
+ let char = B.add_char buf
+ let sub_string = B.add_substring buf
+ let string = B.add_string buf
+ let newline () = B.add_char buf '\n'
+ end)
+
+ module type S = sig
+ type t
+
+ val pre : out -> t -> unit
+ val post : out -> t -> unit
+ end
+
+ type 'a t = (module S with type t = 'a)
+end
+
type t = {
view: view; (** Document view *)
wfl: int; (** Width if flattened *)
@@ -13,6 +43,7 @@ and view =
| Text_sub of string * int * int
| Group of t
| Fill of { sep: t; l: t list }
+ | Wrap : 'a Ext.t * 'a * t -> view
let nil : t = { view = Nil; wfl = 0 }
let newline : t = { view = Newline; wfl = 1 }
@@ -42,6 +73,7 @@ let group d : t =
| Group _ -> d
| _ -> { view = Group d; wfl = d.wfl }
+let wrap ext v d : t = { view = Wrap (ext, v, d); wfl = d.wfl }
let ( ^ ) = append
let text_sub_ s i len : t = { view = Text_sub (s, i, len); wfl = len }
@@ -76,8 +108,6 @@ let text (str : string) : t =
let textpf fmt = Printf.ksprintf text fmt
let textf fmt = Format.kasprintf text fmt
-module B = Buffer
-
module Flatten = struct
let to_buffer buf (self : t) : unit =
let rec loop (d : t) =
@@ -98,7 +128,13 @@ module Flatten = struct
if i > 0 then loop sep;
loop x)
l
+ | Wrap ((module E), v, d) ->
+ let out = Ext.out_of_buf buf in
+ E.pre out v;
+ loop d;
+ E.post out v
in
+
loop self
let to_string self : string =
@@ -145,6 +181,12 @@ module Pretty = struct
n := !n + pp_flatten st x)
l;
!n
+ | Wrap ((module E), v, d) ->
+ let out = Ext.out_of_buf st.buf in
+ E.pre out v;
+ let n = pp_flatten st d in
+ E.post out v;
+ n
(** Does [x] fit in the current line when flattened, given that [k] chars
are already on the line? *)
@@ -194,6 +236,12 @@ module Pretty = struct
) else
pp_rec_top st ~k ~i x kont
| Fill { sep; l } -> pp_fill st ~k ~i sep l kont
+ | Wrap ((module E), v, d) ->
+ let out = Ext.out_of_buf st.buf in
+ E.pre out v;
+ pp_rec_top st ~k ~i d (fun k ->
+ E.post out v;
+ kont k)
and pp_fill st ~k ~i sep l (kont : int -> unit) : unit =
(* [k] is the current offset in the line *)
diff --git a/src/pp/containers_pp.mli b/src/pp/containers_pp.mli
index 30d86aee..5564fa6e 100644
--- a/src/pp/containers_pp.mli
+++ b/src/pp/containers_pp.mli
@@ -55,13 +55,42 @@ val newline : t
val nl : t
(** Alias for {!newline} *)
-(* TODO:
-val ext : pre:(unit -> string) -> post:(unit -> string) -> t -> t
-(** Extension. This is a custom-rendered document.
- TODO: customize how long it is?
- TODO: generic output, not [unit -> string]
- *)
- *)
+(** Extension node.
+
+ In here, we can stuff custom printer nodes. *)
+module Ext : sig
+ module type OUT = sig
+ val char : char -> unit
+ val sub_string : string -> int -> int -> unit
+ val string : string -> unit
+ val newline : unit -> unit
+ end
+
+ type out = (module OUT)
+
+ module type S = sig
+ type t
+
+ val pre : out -> t -> unit
+ val post : out -> t -> unit
+ end
+
+ type 'a t = (module S with type t = 'a)
+end
+
+val wrap : 'a Ext.t -> 'a -> t -> t
+(** [wrap ext v d] wraps [d] with value [v].
+
+ It is a document that has the same
+ shape (and size) as [d], except that additional data will
+ be output when it is rendered.
+
+ Let [(module Ext)] be [ext], and [out]
+ be the output buffer/stream for rendering.
+
+ When this is rendered, first [Ext.pre out v] is called;
+ then [d] is printed; then [Exp.post out v] is called.
+*)
(** {2 Pretty print and rendering} *)
diff --git a/tests/core/t_pp.ml b/tests/core/t_pp.ml
index b782a942..c5ce3309 100644
--- a/tests/core/t_pp.ml
+++ b/tests/core/t_pp.ml
@@ -39,3 +39,22 @@ let () =
[ text "bar"; text "baaz"; sexp_apply "g" [ int 42; int 10 ] ]
in
Pretty.to_string ~width:10 d)
+
+module Ext_coucou : Ext.S with type t = unit = struct
+ type t = unit
+
+ let pre (module O : Ext.OUT) () = O.string ""
+ let post (module O : Ext.OUT) () = O.string ""
+end
+
+let () =
+ eq ~name:"wrap1" ~printer:(spf "%S")
+ "(foo\n bar\n (g 42 10))"
+ (let d =
+ sexp_apply "foo"
+ [
+ text "bar";
+ wrap (module Ext_coucou) () (sexp_apply "g" [ int 42; int 10 ]);
+ ]
+ in
+ Pretty.to_string ~width:10 d)