feat(xml): tree fails in lwt-land directly

2 files changed, 10 insertions, 10 deletions

diff --git a/lib/segment.ml b/lib/segment.ml
index 7023098..b278c50 100644
--- a/lib/segment.ml
+++ b/lib/segment.ml
@@ -1,4 +1,4 @@
-open Lwt.Syntax
+open Lwt.Infix
 open Markup
 open Xml
 
@@ -15,8 +15,4 @@ let next (stream : (signal, async) stream) : (signal, sync) stream Lwt.t =
 
 (** [get stream] is a promise containing a single Xml element of [stream]. *)
 let get (stream : (signal, async) stream) : element Lwt.t =
-  let* signal = next stream in
-  let* el = tree signal
-  in match el with
-     | Some xml -> Lwt.return xml
-     | None -> Lwt.fail (InvalidStanza (signal |> write_xml |> to_string))
+  next stream >>= tree
diff --git a/lib/xml.ml b/lib/xml.ml
index dcb3e79..9c5304e 100644
--- a/lib/xml.ml
+++ b/lib/xml.ml
@@ -1,5 +1,6 @@
 open Lwt.Syntax
 
+exception Malformed
 exception InvalidStanza of string
 
 type element = {
@@ -9,9 +10,9 @@ type element = {
     children : (element, string) Either.t list;
   }
 
-(** [tree s] is a promise to an [element option] representing the XML element inside of
+(** [tree s] is a promise to an [element] representing the XML element inside of
     stream [s], if [s] is a complete element from start to end. *)
-let tree s : element option Lwt.t =
+let tree s : element Lwt.t =
   let element (namespace, name) attributes children =
     Either.Left {
         namespace;
@@ -26,8 +27,11 @@ let tree s : element option Lwt.t =
         children;
       }
   and text ss = Either.Right (String.concat "" ss) in
-  let+ opt_el = Markup_lwt.tree ~text ~element s in
-  Option.bind opt_el (Either.fold ~left:Option.some ~right:(fun _ -> None))
+  let* el = Markup_lwt.tree ~text ~element s
+  in match el with
+     | Some Left el -> Lwt.return el
+     | _ -> Lwt.fail Malformed
+(* Option.bind opt_el (Either.fold ~left:Option.some ~right:(fun _ -> None)) *)
 
 (** [element_to_string element] is a string representation of the underlying XML in
     [element], for debugging purposes.