Add excludedCountries datasource function.

kizitonwose · Aug 2, 2020 · 5cf0fbd · 5cf0fbd
1 parent 2f6fd4c
commit 5cf0fbd
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 10 deletions.
diff --git a/CountryPickerView/CountryPickerView.swift b/CountryPickerView/CountryPickerView.swift
@@ -100,8 +100,8 @@ public class CountryPickerView: NibView {
     internal(set) public var selectedCountry: Country {
         get {
             return _selectedCountry
-                ?? countries.first(where: { $0.code == Locale.current.regionCode })
-                ?? countries.first!
+                ?? usableCountries.first(where: { $0.code == Locale.current.regionCode })
+                ?? usableCountries.first!
         }
         set {
             _selectedCountry = newValue
@@ -198,6 +198,11 @@ public class CountryPickerView: NibView {
         }
         return countries
     }()
+
+    internal var usableCountries: [Country] {
+        let excluded = dataSource?.excludedCountriesInList(in: self) ?? []
+        return countries.filter { return !excluded.contains($0) }
+    }
 }
 
 //MARK: Helper methods

diff --git a/CountryPickerView/CountryPickerViewController.swift b/CountryPickerView/CountryPickerViewController.swift
@@ -45,7 +45,7 @@ extension CountryPickerViewController {
 
     func prepareTableItems()  {
         if !showOnlyPreferredSection {
-            let countriesArray = countryPickerView.countries
+            let countriesArray = countryPickerView.usableCountries
             let locale = dataSource.localeForCountryNameInList
 
             var groupedData = Dictionary<String, [Country]>(grouping: countriesArray) {
@@ -351,4 +351,8 @@ class CountryPickerViewDataSourceInternal: CountryPickerViewDataSource {
     var localeForCountryNameInList: Locale {
         return view.dataSource?.localeForCountryNameInList(in: view) ?? localeForCountryNameInList(in: view)
     }
+
+    var excludedCountries: [Country] {
+        return view.dataSource?.excludedCountriesInList(in: view) ?? preferredCountries(in: view)
+    }
 }
diff --git a/CountryPickerView/Delegate+DataSource.swift b/CountryPickerView/Delegate+DataSource.swift
@@ -79,6 +79,9 @@ public protocol CountryPickerViewDataSource: class {
 
     /// The Locale used to generate the name of the cuntries on the list.
     func localeForCountryNameInList(in countryPickerView: CountryPickerView) -> Locale
+
+    /// An array of countries you wish to exclude from the list of countries.
+    func excludedCountriesInList(in countryPickerView: CountryPickerView) -> [Country]
 }
 
 // MARK:- CountryPickerViewDataSource default implementations
@@ -147,6 +150,10 @@ public extension CountryPickerViewDataSource {
     func localeForCountryNameInList(in countryPickerView: CountryPickerView) -> Locale {
         return Locale.current
     }
+
+    func excludedCountriesInList(in countryPickerView: CountryPickerView) -> [Country] {
+        return []
+    }
 }
 
 

diff --git a/CountryPickerViewDemo/CountryPickerViewDemo/DemoViewController.swift b/CountryPickerViewDemo/CountryPickerViewDemo/DemoViewController.swift
@@ -122,13 +122,7 @@ extension DemoViewController: CountryPickerViewDelegate {
 extension DemoViewController: CountryPickerViewDataSource {
     func preferredCountries(in countryPickerView: CountryPickerView) -> [Country] {
         if countryPickerView.tag == cpvMain.tag && showPreferredCountries.isOn {
-            var countries = [Country]()
-            ["NG", "US", "GB"].forEach { code in
-                if let country = countryPickerView.getCountryByCode(code) {
-                    countries.append(country)
-                }
-            }
-            return countries
+            return ["NG", "US", "GB"].compactMap { countryPickerView.getCountryByCode($0) }
         }
         return []
     }