Handle sub-entities when filtering change sets

Author: kylef

CHANGELOG.md

@@ -6,6 +6,8 @@
 
 - Fixes a bug where `count()` using a predicate may have not counted objects
   that have been modified to come into the predicate.
+- QuerySet subscriptions will now correctly handle abstract and sub entities
+  for managed objects.
 
 ## 0.5.2
 

RxQueryKit/RxQueryKit.swift

@@ -4,26 +4,40 @@ import QueryKit
 import Foundation
 
 
+private func managedObjectMatchesEntity(entityNames: [String]) -> (NSManagedObject) -> Bool {
+  return { object in
+    if let name = object.entity.name {
+      return entityNames.contains(name)
+    }
+
+    return false
+  }
+}
+
+
 /// Extension to QuerySet to provide observables
 extension QuerySet {
+  private var entityNames: [String] {
+    if let entity = NSEntityDescription.entityForName(entityName, inManagedObjectContext: context) {
+      return [entityName] + entity.subentitiesByName.keys
+    }
+
+    return [entityName]
+  }
+
   /// Performs a query for all objects matching the set predicate ordered by any set sort descriptors.
   /// Emits a value with an array of all objects when the managed object context is changed.
   public func objects() throws -> Observable<[ModelType]> {
     let initial = try self.array()
+    let filter = managedObjectMatchesEntity(entityNames)
 
     return Observable.create { observer in
       observer.on(.Next(initial))
 
       return self.context.qk_objectsDidChange().subscribeNext { notification in
-        let insertedObjects = notification.insertedObjects.filter {
-          $0.entity.name == self.entityName
-        }
-        let updatedObjects = notification.updatedObjects.filter {
-          $0.entity.name == self.entityName
-        }
-        let deletedObjects = notification.deletedObjects.filter {
-          $0.entity.name == self.entityName
-        }
+        let insertedObjects = notification.insertedObjects.filter(filter)
+        let updatedObjects = notification.updatedObjects.filter(filter)
+        let deletedObjects = notification.deletedObjects.filter(filter)
 
         if insertedObjects.isEmpty && updatedObjects.isEmpty && deletedObjects.isEmpty {
           return
@@ -43,14 +57,13 @@ extension QuerySet {
   /// Emits an Int containing the amount of objects matching the predicate and updates when the managed object context is changed.
   public func count() throws -> Observable<Int> {
     var count: Int = try self.count()
+    let filter = managedObjectMatchesEntity(entityNames)
 
     return Observable.create { observer in
       observer.on(.Next(count))
 
       return self.context.qk_objectsDidChange().subscribeNext { notification in
-        let updatedObjects = notification.updatedObjects.filter {
-          $0.entity.name == self.entityName
-        }
+        let updatedObjects = notification.updatedObjects.filter(filter)
 
         if !updatedObjects.isEmpty && self.predicate != nil {
           do {
@@ -62,12 +75,8 @@ extension QuerySet {
           return
         }
 
-        let insertedObjects = notification.insertedObjects.filter {
-          $0.entity.name == self.entityName
-        }
-        let deletedObjects = notification.deletedObjects.filter {
-          $0.entity.name == self.entityName
-        }
+        let insertedObjects = notification.insertedObjects.filter(filter)
+        let deletedObjects = notification.deletedObjects.filter(filter)
 
         var delta = 0