feat(stream): return a stream of signals when using get

1 files changed, 10 insertions, 7 deletions

diff --git a/lib/stream.ml b/lib/stream.ml
index b26d02f..dff2b0e 100644
--- a/lib/stream.ml
+++ b/lib/stream.ml
@@ -1,16 +1,19 @@
 open Lwt.Syntax
+open Markup
 
 exception ClosedStream
 exception InvalidStanza of string
 
-let get stream =
-  (** [get stream] returns a promise containing an element of [stream].
+let get (stream : (signal, async) stream) : (signal, sync) stream Lwt.t =
+  (** [stanza stream] is a promise containing a full stanza of the fragments of [stream].*)
 
-      If [stream] is closed, it throws a ClosedStream Lwt exception. *)
-  let* stanza = Lwt_stream.get stream
-  in match stanza with
-     | Some stanza -> Lwt.return stanza
-     | None -> Lwt.fail ClosedStream
+  let traverse_stanza depth fragment =
+    let depth = match fragment with
+      | `Start_element _ -> depth + 1
+      | `End_element -> depth - 1
+      | _ -> depth
+    in ([fragment], if depth = 0 then None else Some depth)
+  in transform traverse_stanza 0 stream |> Markup_lwt.load
 
 let start domain : Portal.t Lwt.t =
   (** [start domain] is a promise containing a Portal (stream * push) connected to the XMPP server [domain].