From 1556ff01d85e8e756f9a6c40cc531fe5cad461a1 Mon Sep 17 00:00:00 2001
From: David Beck <code@thinkultimate.com>
Date: Mon, 30 Mar 2015 15:57:14 -0700
Subject: [PATCH] Fixed mediaTypes filter

---
 .../TNKViewController.m                       |  1 +
 Pod/Classes/TNKImagePickerController.m        | 28 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/Example/TNKImagePickerController/TNKViewController.m b/Example/TNKImagePickerController/TNKViewController.m
index 5bac8d2..7c1b560 100644
--- 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
index ef92f00..c5f0a96 100644
--- 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;