Add an option to STPPaymentConfiguration to filter the country list

Stripe/PublicHeaders/STPPaymentConfiguration.h

@@ -77,6 +77,15 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property (nonatomic, assign, readwrite) STPShippingType shippingType;
 
+/**
+ The set of countries supported when entering an address. This property accepts
+ a set of ISO 2-character country codes.
+
+ The default value is nil, which will display all known countries. Setting this
+ property will limit the available countries to your selected set.
+ */
+@property (nonatomic, copy, nullable, readwrite) NSSet<NSString *> *availableCountries;
+
 /**
  The name of your company, for displaying to the user during payment flows. For 
  example, when using Apple Pay, the payment sheet's final line item will read

Stripe/STPAddCardViewController.m

@@ -93,7 +93,7 @@ - (void)commonInitWithConfiguration:(STPPaymentConfiguration *)configuration {
     _shippingAddress = nil;
     _hasUsedShippingAddress = NO;
     _apiClient = [[STPAPIClient alloc] initWithConfiguration:configuration];
-    _addressViewModel = [[STPAddressViewModel alloc] initWithRequiredBillingFields:configuration.requiredBillingAddressFields];
+    _addressViewModel = [[STPAddressViewModel alloc] initWithRequiredBillingFields:configuration.requiredBillingAddressFields availableCountries:configuration.availableCountries];
     _addressViewModel.delegate = self;
 
     self.title = STPLocalizedString(@"Add a Card", @"Title for Add a Card view");

Stripe/STPAddressFieldTableViewCell.h

@@ -34,6 +34,7 @@ typedef NS_ENUM(NSInteger, STPAddressFieldType) {
 - (void)addressFieldTableViewCellDidReturn:(STPAddressFieldTableViewCell *)cell;
 - (void)addressFieldTableViewCellDidEndEditing:(STPAddressFieldTableViewCell *)cell;
 @property (nonatomic, copy) NSString *addressFieldTableViewCountryCode;
+@property (nonatomic, copy) NSSet<NSString *> *availableCountries;
 
 @end
 

Stripe/STPAddressFieldTableViewCell.m

@@ -69,9 +69,18 @@ - (instancetype)initWithType:(STPAddressFieldType)type
         _inputAccessoryToolbar = toolbar;
 
         NSString *countryCode = [[NSLocale autoupdatingCurrentLocale] objectForKey:NSLocaleCountryCode];
-        NSMutableArray *otherCountryCodes = [[NSLocale ISOCountryCodes] mutableCopy];
+        NSMutableArray *otherCountryCodes = [[self.delegate.availableCountries allObjects] mutableCopy];
+        if (otherCountryCodes == nil) {            
+            otherCountryCodes = [[NSLocale ISOCountryCodes] mutableCopy];
+        }
+        if ([otherCountryCodes containsObject:countryCode]) {
+            // Remove the current country code to re-add it once we sort the list.
+            [otherCountryCodes removeObject:countryCode];
+        } else {
+            // If it isn't in the list (if we've been configured to not show that country), don't re-add it.
+            countryCode = nil;
+        }
         NSLocale *locale = [NSLocale currentLocale];
-        [otherCountryCodes removeObject:countryCode];
         [otherCountryCodes sortUsingComparator:^NSComparisonResult(NSString *code1, NSString *code2) {
             NSString *localeID1 = [NSLocale localeIdentifierFromComponents:@{NSLocaleCountryCode: code1}];
             NSString *localeID2 = [NSLocale localeIdentifierFromComponents:@{NSLocaleCountryCode: code2}];
@@ -179,6 +188,11 @@ - (void)updateTextFieldsAndCaptions {
             self.textField.keyboardType = UIKeyboardTypeDefault;
             // Don't set textContentType for Country, because we don't want iOS to skip the UIPickerView for input
             self.textField.inputView = self.countryPickerView;
+
+            // If we're being set directly to a country we don't allow, add it to the allowed list
+            if (![self.countryCodes containsObject:self.contents] && [[NSLocale ISOCountryCodes] containsObject:self.contents]) {
+                self.countryCodes = [self.countryCodes arrayByAddingObject:self.contents];
+            }
             NSInteger index = [self.countryCodes indexOfObject:self.contents];
             if (index == NSNotFound) {
                 self.textField.text = @"";

Stripe/STPAddressViewModel.h

@@ -25,10 +25,13 @@
 @property (nonatomic, readonly) NSArray<STPAddressFieldTableViewCell *> *addressCells;
 @property (nonatomic, weak) id<STPAddressViewModelDelegate>delegate;
 @property (nonatomic) STPAddress *address;
+@property (nonatomic, copy, readwrite) NSSet<NSString *> *availableCountries;
 @property (nonatomic, readonly) BOOL isValid;
 
 - (instancetype)initWithRequiredBillingFields:(STPBillingAddressFields)requiredBillingAddressFields;
 - (instancetype)initWithRequiredShippingFields:(NSSet<STPContactField> *)requiredShippingAddressFields;
+- (instancetype)initWithRequiredBillingFields:(STPBillingAddressFields)requiredBillingAddressFields availableCountries:(NSSet<NSString *> *)availableCountries;
+- (instancetype)initWithRequiredShippingFields:(NSSet<STPContactField> *)requiredShippingAddressFields availableCountries:(NSSet<NSString *> *)availableCountries;
 - (STPAddressFieldTableViewCell *)cellAtIndex:(NSInteger)index;
 
 @end

Stripe/STPAddressViewModel.m

@@ -26,11 +26,13 @@ @interface STPAddressViewModel()<STPAddressFieldTableViewCellDelegate>
 @implementation STPAddressViewModel
 
 @synthesize addressFieldTableViewCountryCode = _addressFieldTableViewCountryCode;
+@synthesize availableCountries = _availableCountries;
 
-- (instancetype)initWithRequiredBillingFields:(STPBillingAddressFields)requiredBillingAddressFields {
+- (instancetype)initWithRequiredBillingFields:(STPBillingAddressFields)requiredBillingAddressFields availableCountries:(NSSet<NSString *> *)availableCountries {
     self = [super init];
     if (self) {
         _isBillingAddress = YES;
+        _availableCountries = [availableCountries copy];
         _requiredBillingAddressFields = requiredBillingAddressFields;
         switch (requiredBillingAddressFields) {
             case STPBillingAddressFieldsNone:
@@ -63,10 +65,11 @@ - (instancetype)initWithRequiredBillingFields:(STPBillingAddressFields)requiredB
     return self;
 }
 
-- (instancetype)initWithRequiredShippingFields:(NSSet<STPContactField> *)requiredShippingAddressFields {
+- (instancetype)initWithRequiredShippingFields:(NSSet<STPContactField> *)requiredShippingAddressFields availableCountries:(NSSet<NSString *> *)availableCountries {
     self = [super init];
     if (self) {
         _isBillingAddress = NO;
+        _availableCountries = [availableCountries copy];
         _requiredShippingAddressFields = requiredShippingAddressFields;
         NSMutableArray *cells = [NSMutableArray new];
         if ([requiredShippingAddressFields containsObject:STPContactFieldName]) {
@@ -103,6 +106,13 @@ - (instancetype)initWithRequiredShippingFields:(NSSet<STPContactField> *)require
     return self;
 }
 
+- (instancetype)initWithRequiredBillingFields:(STPBillingAddressFields)requiredBillingAddressFields {
+    return [self initWithRequiredBillingFields:requiredBillingAddressFields availableCountries:nil];
+}
+- (instancetype)initWithRequiredShippingFields:(NSSet<STPContactField> *)requiredShippingAddressFields {
+    return [self initWithRequiredShippingFields:requiredShippingAddressFields availableCountries:nil];
+}
+
 - (void)commonInit {
     _addressFieldTableViewCountryCode = [[NSLocale autoupdatingCurrentLocale] objectForKey:NSLocaleCountryCode];
     [self updatePostalCodeCellIfNecessary];

Stripe/STPPaymentConfiguration.m

@@ -120,6 +120,7 @@ - (NSString *)description {
                        [NSString stringWithFormat:@"requiredShippingAddressFields = %@", requiredShippingAddressFieldsDescription],
                        [NSString stringWithFormat:@"verifyPrefilledShippingAddress = %@", (self.verifyPrefilledShippingAddress) ? @"YES" : @"NO"],
                        [NSString stringWithFormat:@"shippingType = %@", shippingTypeDescription],
+                       [NSString stringWithFormat:@"availableCountries = %@", _availableCountries],
 
                        // Additional configuration
                        [NSString stringWithFormat:@"companyName = %@", self.companyName],
@@ -143,6 +144,7 @@ - (id)copyWithZone:(__unused NSZone *)zone {
     copy.companyName = self.companyName;
     copy.appleMerchantIdentifier = self.appleMerchantIdentifier;
     copy.canDeletePaymentOptions = self.canDeletePaymentOptions;
+    copy.availableCountries = self.availableCountries;
     return copy;
 }
 

Stripe/STPShippingAddressViewController.m

@@ -87,7 +87,7 @@ - (instancetype)initWithConfiguration:(STPPaymentConfiguration *)configuration
         _selectedShippingMethod = selectedShippingMethod;
         _billingAddress = prefilledInformation.billingAddress;
         _hasUsedBillingAddress = NO;
-        _addressViewModel = [[STPAddressViewModel alloc] initWithRequiredShippingFields:configuration.requiredShippingAddressFields];
+        _addressViewModel = [[STPAddressViewModel alloc] initWithRequiredShippingFields:configuration.requiredShippingAddressFields availableCountries:configuration.availableCountries];
         _addressViewModel.delegate = self;
         if (shippingAddress != nil) {
             _addressViewModel.address = shippingAddress;

Tests/Tests/STPPaymentConfigurationTest.m

@@ -107,6 +107,7 @@ - (void)testCopyWithZone {
     paymentConfigurationA.requiredBillingAddressFields = STPBillingAddressFieldsFull;
     paymentConfigurationA.requiredShippingAddressFields = allFields;
     paymentConfigurationA.verifyPrefilledShippingAddress = NO;
+    paymentConfigurationA.availableCountries = [NSSet setWithArray:@[@"US", @"CA", @"BT"]];
     paymentConfigurationA.shippingType = STPShippingTypeDelivery;
     paymentConfigurationA.companyName = @"companyName";
     paymentConfigurationA.appleMerchantIdentifier = @"appleMerchantIdentifier";
@@ -123,6 +124,8 @@ - (void)testCopyWithZone {
     XCTAssertEqual(paymentConfigurationB.shippingType, STPShippingTypeDelivery);
     XCTAssertEqualObjects(paymentConfigurationB.companyName, @"companyName");
     XCTAssertEqualObjects(paymentConfigurationB.appleMerchantIdentifier, @"appleMerchantIdentifier");
+    NSSet *availableCountries = [NSSet setWithArray:@[@"US", @"CA", @"BT"]];
+    XCTAssertEqualObjects(paymentConfigurationB.availableCountries, availableCountries);
     XCTAssertEqual(paymentConfigurationA.canDeletePaymentOptions, paymentConfigurationB.canDeletePaymentOptions);
 }
 

Tests/Tests/STPShippingAddressViewControllerTest.m

@@ -46,6 +46,38 @@ - (void)testPrefilledBillingAddress_removeAddress {
     XCTAssertNoThrow([sut viewDidLoad]);
 }
 
+- (void)testPrefilledBillingAddress_addAddressWithLimitedCountries {
+    [NSLocale stp_setCurrentLocale:[NSLocale localeWithLocaleIdentifier:@"en_ZW"]];
+    // Zimbabwe does not require zip codes, while the default locale for tests (US) does
+    // Sanity checks
+    XCTAssertFalse([STPPostalCodeValidator postalCodeIsRequiredForCountryCode:@"ZW"]);
+    XCTAssertTrue([STPPostalCodeValidator postalCodeIsRequiredForCountryCode:@"US"]);
+    STPPaymentConfiguration *config = [STPFixtures paymentConfiguration];
+    config.requiredShippingAddressFields = [NSSet setWithObject:STPContactFieldPostalAddress];
+    config.availableCountries = [[NSSet alloc] initWithArray:@[@"CA", @"BT"]];
+
+    STPAddress *address = [STPAddress new];
+    address.name = @"John Smith Doe";
+    address.phone = @"8885551212";
+    address.email = @"[email protected]";
+    address.line1 = @"55 John St";
+    address.city = @"New York";
+    address.state = @"NY";
+    address.postalCode = @"10002";
+    address.country = @"US";
+
+    STPShippingAddressViewController *sut = [[STPShippingAddressViewController alloc] initWithConfiguration:config
+                                                                                                      theme:[STPTheme defaultTheme]
+                                                                                                   currency:nil
+                                                                                            shippingAddress:address
+                                                                                     selectedShippingMethod:nil
+                                                                                       prefilledInformation:nil];
+
+    XCTAssertNoThrow([sut loadView]);
+    XCTAssertNoThrow([sut viewDidLoad]);
+    [NSLocale stp_resetCurrentLocale];
+}
+
 - (void)testPrefilledBillingAddress_addAddress {
     [NSLocale stp_setCurrentLocale:[NSLocale localeWithLocaleIdentifier:@"en_ZW"]];
     // Zimbabwe does not require zip codes, while the default locale for tests (US) does