AsyncSeq.traverseOptionAsync / traverseChoiceAsync

Author: eulerfx

src/FSharpx.Async/AsyncSeq.fs

@@ -7,6 +7,7 @@ namespace FSharpx.Control
 open System
 open System.Threading
 open System.IO
+open FSharpx.Control.Utils
 
 // ----------------------------------------------------------------------------
 
@@ -474,6 +475,34 @@ module AsyncSeq =
   let inline zapp (fs:AsyncSeq<'a -> 'b>) (s:AsyncSeq<'a>) : AsyncSeq<'b> =
     zipWith ((|>)) s fs
 
+  /// Traverses an async sequence an applies to specified function such that if None is returned the traversal short-circuits
+  /// and None is returned as the result. Otherwise, the entire sequence is traversed and the result returned as Some.
+  let rec traverseOptionAsync (f:'a -> Async<'b option>) (s:AsyncSeq<'a>) : Async<AsyncSeq<'b> option> = async {
+    let! s = s
+    match s with
+    | Nil -> return Some (Nil |> async.Return)
+    | Cons(a,tl) ->
+      let! b = f a
+      match b with
+      | Some b -> 
+        return! traverseOptionAsync f tl |> Async.map (Option.map (fun tl -> Cons(b, tl) |> async.Return))
+      | None -> 
+        return None }
+
+  /// Traverses an async sequence an applies to specified function such that if Choice2Of2 is returned the traversal short-circuits
+  /// and Choice2Of2 is returned as the result. Otherwise, the entire sequence is traversed and the result returned as Choice1Of2.
+  let rec traverseChoiceAsync (f:'a -> Async<Choice<'b, 'e>>) (s:AsyncSeq<'a>) : Async<Choice<AsyncSeq<'b>, 'e>> = async {
+    let! s = s
+    match s with
+    | Nil -> return Choice1Of2 (Nil |> async.Return)
+    | Cons(a,tl) ->
+      let! b = f a
+      match b with
+      | Choice1Of2 b -> 
+        return! traverseChoiceAsync f tl |> Async.map (Choice.mapl (fun tl -> Cons(b, tl) |> async.Return))
+      | Choice2Of2 e -> 
+        return Choice2Of2 e }
+
   /// Returns an async computation which iterates the async sequence for
   /// its side-effects.
   let inline runAsync (s:AsyncSeq<unit>) : Async<unit> =

tests/FSharpx.Async.Tests/AsyncSeqTests.fs

@@ -209,3 +209,31 @@ let ``AsyncSeq.replicate``() =
   let expected = List.replicate c x |> AsyncSeq.ofSeq
   Assert.True(EQ expected actual)
 
+[<Test>]
+let ``AsyncSeq.traverseOptionAsync``() =
+  let seen = ResizeArray<_>()
+  let s = [1;2;3;4;5] |> AsyncSeq.ofSeq
+  let f i =
+    seen.Add i
+    if i < 2 then Some i |> async.Return
+    else None |> async.Return
+  let r = AsyncSeq.traverseOptionAsync f s |> Async.RunSynchronously
+  match r with
+  | Some _ -> Assert.Fail()
+  | None -> Assert.True(([1;2] = (seen |> List.ofSeq)))
+
+[<Test>]
+let ``AsyncSeq.traverseChoiceAsync``() =
+  let seen = ResizeArray<_>()
+  let s = [1;2;3;4;5] |> AsyncSeq.ofSeq
+  let f i =
+    seen.Add i
+    if i < 2 then Choice1Of2 i |> async.Return
+    else Choice2Of2 "oh no" |> async.Return
+  let r = AsyncSeq.traverseChoiceAsync f s |> Async.RunSynchronously
+  match r with
+  | Choice1Of2 _ -> Assert.Fail()
+  | Choice2Of2 e -> 
+    Assert.AreEqual("oh no", e)
+    Assert.True(([1;2] = (seen |> List.ofSeq)))
+