diff --git a/CountryPickerView/CountryPickerView.swift b/CountryPickerView/CountryPickerView.swift index f811e6d..ae38fb0 100644 --- 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 index e2b7166..1d02692 100644 --- 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(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 index 3ce33fd..50bf44d 100644 --- 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 index 2987e71..0c0119c 100644 --- 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 [] }