Fixed mediaTypes filter

davbeck · Mar 30, 2015 · 1556ff0 · 1556ff0
1 parent 434a42d
commit 1556ff0
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/Example/TNKImagePickerController/TNKViewController.m b/Example/TNKImagePickerController/TNKViewController.m
@@ -32,6 +32,7 @@ - (void)viewDidLoad
 - (IBAction)pickPhotos:(id)sender
 {
     TNKImagePickerController *picker = [[TNKImagePickerController alloc] init];
+    picker.mediaTypes = @[ (id)kUTTypeImage ];
 
     UINavigationController *navigationController = [[UINavigationController alloc] initWithRootViewController:picker];
     navigationController.toolbarHidden = NO;

diff --git a/Pod/Classes/TNKImagePickerController.m b/Pod/Classes/TNKImagePickerController.m
@@ -75,6 +75,21 @@ - (void)deselectAsset:(PHAsset *)asset {
     [self removeSelectedAssets:[NSOrderedSet orderedSetWithObject:asset]];
 }
 
+- (NSArray *)_assetMediaTypes {
+    NSMutableArray *assetMediaTypes = [NSMutableArray new];
+    if ([self.mediaTypes containsObject:(id)kUTTypeImage]) {
+        [assetMediaTypes addObject:@(PHAssetMediaTypeImage)];
+    }
+    if ([self.mediaTypes containsObject:(id)kUTTypeVideo]) {
+        [assetMediaTypes addObject:@(PHAssetMediaTypeVideo)];
+    }
+    if ([self.mediaTypes containsObject:(id)kUTTypeAudio]) {
+        [assetMediaTypes addObject:@(PHAssetMediaTypeAudio)];
+    }
+
+    return assetMediaTypes;
+}
+
 - (void)setAssetCollection:(PHAssetCollection *)assetCollection {
     _assetCollection = assetCollection;
 
@@ -94,7 +109,10 @@ - (void)_updateForAssetCollection
 
 
     if (_assetCollection != nil) {
-        _fetchResult = [PHAsset fetchAssetsInAssetCollection:_assetCollection options:nil];
+        PHFetchOptions *options = [[PHFetchOptions alloc] init];
+        options.predicate = [NSPredicate predicateWithFormat:@"mediaType IN %@", [self _assetMediaTypes]];
+
+        _fetchResult = [PHAsset fetchAssetsInAssetCollection:_assetCollection options:options];
         _moments = nil;
     } else {
         _fetchResult = nil;
@@ -543,6 +561,7 @@ - (PHFetchResult *)_assetsForMoment:(PHAssetCollection *)collection
     if (result == nil) {
         PHFetchOptions *options = [PHFetchOptions new];
         options.includeAllBurstAssets = NO;
+        options.predicate = [NSPredicate predicateWithFormat:@"mediaType IN %@", [self _assetMediaTypes]];
         result = [PHAsset fetchAssetsInAssetCollection:collection options:options];
         [_momentCache setObject:result forKey:collection.localIdentifier];
     }
@@ -611,6 +630,13 @@ - (UICollectionReusableView *)collectionView:(UICollectionView *)collectionView
 
 - (CGSize)collectionView:(UICollectionView *)collectionView layout:(UICollectionViewLayout *)collectionViewLayout referenceSizeForHeaderInSection:(NSInteger)section {
     if (_moments != nil) {
+        PHAssetCollection *collection = _moments[section];
+
+        PHFetchResult *fetchResult = [self _assetsForMoment:collection];
+        if (fetchResult.count == 0) {
+            return CGSizeZero;
+        }
+
         return CGSizeMake(collectionView.bounds.size.width, 44.0);
     } else {
         return CGSizeZero;