refactor: use makeStream for creating AsyncStream with continuation

aws-amplify · May 17, 2024 · d5eec5f · d5eec5f
1 parent 5798204
commit d5eec5f
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 21 deletions.
diff --git a/Amplify/Core/Support/AmplifyAsyncSequence.swift b/Amplify/Core/Support/AmplifyAsyncSequence.swift
@@ -11,18 +11,16 @@ public typealias WeakAmplifyAsyncSequenceRef<Element> = WeakRef<AmplifyAsyncSequ
 
 public class AmplifyAsyncSequence<Element: Sendable>: AsyncSequence, Cancellable {
     public typealias Iterator = AsyncStream<Element>.Iterator
-    private var asyncStream: AsyncStream<Element>! = nil
-    private var continuation: AsyncStream<Element>.Continuation! = nil
+    private var asyncStream: AsyncStream<Element>
+    private var continuation: AsyncStream<Element>.Continuation
     private var parent: Cancellable?
 
     public private(set) var isCancelled: Bool = false
 
     public init(parent: Cancellable? = nil,
                 bufferingPolicy: AsyncStream<Element>.Continuation.BufferingPolicy = .unbounded) {
         self.parent = parent
-        asyncStream = AsyncStream<Element>(Element.self, bufferingPolicy: bufferingPolicy) { continuation in
-            self.continuation = continuation
-        }
+        (asyncStream, continuation) = AsyncStream.makeStream(of: Element.self, bufferingPolicy: bufferingPolicy)
     }
 
     public func makeAsyncIterator() -> Iterator {

diff --git a/Amplify/Core/Support/AmplifyAsyncThrowingSequence.swift b/Amplify/Core/Support/AmplifyAsyncThrowingSequence.swift
@@ -11,18 +11,16 @@ public typealias WeakAmplifyAsyncThrowingSequenceRef<Element> = WeakRef<AmplifyA
 
 public class AmplifyAsyncThrowingSequence<Element: Sendable>: AsyncSequence, Cancellable {
     public typealias Iterator = AsyncThrowingStream<Element, Error>.Iterator
-    private var asyncStream: AsyncThrowingStream<Element, Error>! = nil
-    private var continuation: AsyncThrowingStream<Element, Error>.Continuation! = nil
+    private var asyncStream: AsyncThrowingStream<Element, Error>
+    private var continuation: AsyncThrowingStream<Element, Error>.Continuation
     private var parent: Cancellable?
 
     public private(set) var isCancelled: Bool = false
 
     public init(parent: Cancellable? = nil,
                 bufferingPolicy: AsyncThrowingStream<Element, Error>.Continuation.BufferingPolicy = .unbounded) {
         self.parent = parent
-        asyncStream = AsyncThrowingStream(Element.self, bufferingPolicy: bufferingPolicy, { continuation in
-            self.continuation = continuation
-        })
+        (asyncStream, continuation) = AsyncThrowingStream.makeStream(of: Element.self, bufferingPolicy: bufferingPolicy)
     }
 
     public func makeAsyncIterator() -> Iterator {

diff --git a/Amplify/Core/Support/TaskQueue.swift b/Amplify/Core/Support/TaskQueue.swift
@@ -13,9 +13,8 @@ public class TaskQueue<Success> {
     private var streamContinuation: AsyncStream<Block>.Continuation!
 
     public init() {
-        let stream = AsyncStream<Block>.init { continuation in
-            streamContinuation = continuation
-        }
+        let (stream, continuation) = AsyncStream.makeStream(of: Block.self)
+        self.streamContinuation = continuation
 
         Task {
             for await block in stream {

diff --git a/...nitoAuthPlugin/StateMachine/hierarchical-state-machine-swift/CancellableAsyncStream.swift b/...nitoAuthPlugin/StateMachine/hierarchical-state-machine-swift/CancellableAsyncStream.swift
@@ -23,14 +23,14 @@ class CancellableAsyncStream<Element>: AsyncSequence {
     }
 
     convenience init(with publisher: AnyPublisher<Element, Never>) {
-        var cancellable: AnyCancellable?
-        self.init(asyncStream: AsyncStream { continuation in
-            cancellable = publisher.sink { _ in
-                continuation.finish()
-            } receiveValue: {
-                continuation.yield($0)
-            }
-        }, cancellable: cancellable)
+        let (asyncStream, contiuation) = AsyncStream.makeStream(of: Element.self)
+        let cancellable = publisher.sink { _ in
+            contiuation.finish()
+        } receiveValue: {
+            contiuation.yield($0)
+        }
+
+        self.init(asyncStream: asyncStream, cancellable: cancellable)
     }
 
     func makeAsyncIterator() -> AsyncStream<Element>.AsyncIterator {