Add completion block to apply function

Author: ra1028

Sources/AppKit/CocoaCollectionViewDiffableDataSource.swift

@@ -33,14 +33,18 @@ open class CocoaCollectionViewDiffableDataSource<SectionIdentifierType: Hashable
     ///   - snapshot: A snapshot object to be applied to data model.
     ///   - animatingDifferences: A Boolean value indicating whether to update with
     ///                           diffing animation.
+    ///   - completion: An optional completion block which is called when the complete
+    ///                 performing updates.
     public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true) {
         core.apply(
             snapshot,
             view: collectionView,
             animatingDifferences: animatingDifferences,
             performUpdates: { collectionView, changeset, setSections in
                 collectionView.reload(using: changeset, setData: setSections)
-        })
+        },
+            completion: completion
+        )
     }
 
     /// Returns a new snapshot object of current state.

Sources/Internal/DiffableDataSourceCore.swift

@@ -13,7 +13,8 @@ final class DiffableDataSourceCore<SectionIdentifierType: Hashable, ItemIdentifi
         _ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>,
         view: View?,
         animatingDifferences: Bool,
-        performUpdates: @escaping (View, StagedChangeset<[Section]>, @escaping ([Section]) -> Void) -> Void
+        performUpdates: @escaping (View, StagedChangeset<[Section]>, @escaping ([Section]) -> Void) -> Void,
+        completion: (() -> Void)?
         ) {
         dispatcher.dispatch { [weak self] in
             guard let self = self else {
@@ -35,15 +36,18 @@ final class DiffableDataSourceCore<SectionIdentifierType: Hashable, ItemIdentifi
                 }
             }
 
+            CATransaction.begin()
+            CATransaction.setCompletionBlock(completion)
+
             if animatingDifferences {
                 performDiffingUpdates()
             }
             else {
-                CATransaction.begin()
                 CATransaction.setDisableActions(true)
                 performDiffingUpdates()
-                CATransaction.commit()
             }
+
+            CATransaction.commit()
         }
     }
 

Sources/UIKit/CollectionViewDiffableDataSource.swift

@@ -39,14 +39,18 @@ open class CollectionViewDiffableDataSource<SectionIdentifierType: Hashable, Ite
     ///   - snapshot: A snapshot object to be applied to data model.
     ///   - animatingDifferences: A Boolean value indicating whether to update with
     ///                           diffing animation.
-    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true) {
+    ///   - completion: An optional completion block which is called when the complete
+    ///                 performing updates.
+    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true, completion: (() -> Void)? = nil) {
         core.apply(
             snapshot,
             view: collectionView,
             animatingDifferences: animatingDifferences,
             performUpdates: { collectionView, changeset, setSections in
                 collectionView.reload(using: changeset, setData: setSections)
-        })
+        },
+            completion: completion
+        )
     }
 
     /// Returns a new snapshot object of current state.

Sources/UIKit/TableViewDiffableDataSource.swift

@@ -36,14 +36,18 @@ open class TableViewDiffableDataSource<SectionIdentifierType: Hashable, ItemIden
     ///   - snapshot: A snapshot object to be applied to data model.
     ///   - animatingDifferences: A Boolean value indicating whether to update with
     ///                           diffing animation.
-    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true) {
+    ///   - completion: An optional completion block which is called when the complete
+    ///                 performing updates.
+    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true, completion: (() -> Void)? = nil) {
         core.apply(
             snapshot,
             view: tableView,
             animatingDifferences: animatingDifferences,
             performUpdates: { tableView, changeset, setSections in
                 tableView.reload(using: changeset, with: self.defaultRowAnimation, setData: setSections)
-        })
+        },
+            completion: completion
+        )
     }
 
     /// Returns a new snapshot object of current state.