From 4100525890910d431d5e86c7f9c5f067fca9a9c1 Mon Sep 17 00:00:00 2001
From: Ben Guo <bg@stripe.com>
Date: Fri, 21 Oct 2016 14:41:43 -0400
Subject: [PATCH] dismissWithCompletion

---
 .../CheckoutViewController.swift                |  2 +-
 .../STPPaymentMethodsViewController.h           |  7 +++++++
 .../STPShippingAddressViewController.h          |  4 ++--
 Stripe/STPPaymentMethodsViewController.m        | 17 +++++++++++++++++
 Stripe/STPShippingAddressViewController.m       | 16 ++++++++++++++--
 5 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/Example/Stripe iOS Example (Simple)/CheckoutViewController.swift b/Example/Stripe iOS Example (Simple)/CheckoutViewController.swift
index 0da2503663..4fa543defc 100644
--- a/Example/Stripe iOS Example (Simple)/CheckoutViewController.swift	
+++ b/Example/Stripe iOS Example (Simple)/CheckoutViewController.swift	
@@ -256,7 +256,7 @@ class CheckoutViewController: UIViewController, STPPaymentContextDelegate, STPSh
         else {
             self.shippingRow.detail = "Enter \(self.shippingString) Info"
         }
-        self.shippingVC.dismiss(withHostViewController: self)
+        self.shippingVC.dismiss(completion: nil)
     }
 
 }
diff --git a/Stripe/PublicHeaders/STPPaymentMethodsViewController.h b/Stripe/PublicHeaders/STPPaymentMethodsViewController.h
index 1bbd7f93e2..e63445aac7 100644
--- a/Stripe/PublicHeaders/STPPaymentMethodsViewController.h
+++ b/Stripe/PublicHeaders/STPPaymentMethodsViewController.h
@@ -53,6 +53,13 @@ NS_ASSUME_NONNULL_BEGIN
 */
 @property(nonatomic)STPUserInformation *prefilledInformation;
 
+/**
+ *  If you're pushing `STPPaymentMethodsViewController` onto an existing `UINavigationController`'s stack, you should use this method to dismiss it, since it may have pushed an additional add card view controller onto the navigation controller's stack.
+ *
+ *  @param completion The callback to run after the view controller is dismissed. You may specify nil for this parameter.
+ */
+- (void)dismissWithCompletion:(nullable STPVoidBlock)completion;
+
 @end
 
 /**
diff --git a/Stripe/PublicHeaders/STPShippingAddressViewController.h b/Stripe/PublicHeaders/STPShippingAddressViewController.h
index 3f225faf8a..ad278ce34e 100644
--- a/Stripe/PublicHeaders/STPShippingAddressViewController.h
+++ b/Stripe/PublicHeaders/STPShippingAddressViewController.h
@@ -47,9 +47,9 @@ NS_ASSUME_NONNULL_BEGIN
 /**
  *  If you're pushing `STPShippingAddressViewController` onto an existing `UINavigationController`'s stack, you should use this method to dismiss it, since it may have pushed an additional shipping method view controller onto the navigation controller's stack.
  *
- *  @param hostViewController the view controller that the shipping address view controller was pushed onto.
+ *  @param completion The callback to run after the view controller is dismissed. You may specify nil for this parameter.
  */
-- (void)dismissWithHostViewController:(UIViewController *)hostViewController;
+- (void)dismissWithCompletion:(nullable STPVoidBlock)completion;
 
 @end
 
diff --git a/Stripe/STPPaymentMethodsViewController.m b/Stripe/STPPaymentMethodsViewController.m
index 2a648e1ffb..039ea29c65 100644
--- a/Stripe/STPPaymentMethodsViewController.m
+++ b/Stripe/STPPaymentMethodsViewController.m
@@ -27,6 +27,7 @@
 #import "STPWeakStrongMacros.h"
 #import "STPLocalizationUtils.h"
 #import "STPDispatchFunctions.h"
+#import "UINavigationController+Stripe_Completion.h"
 
 @interface STPPaymentMethodsViewController()<STPPaymentMethodsInternalViewControllerDelegate, STPAddCardViewControllerDelegate>
 
@@ -210,6 +211,22 @@
     [self internalViewControllerDidCreateToken:token completion:completion];
 }
 
+- (void)dismissWithCompletion:(STPVoidBlock)completion {
+    if ([self stp_isAtRootOfNavigationController]) {
+        [self.presentingViewController dismissViewControllerAnimated:YES completion:completion];
+    }
+    else {
+        UIViewController *previous = self.navigationController.viewControllers.firstObject;
+        for (UIViewController *viewController in self.navigationController.viewControllers) {
+            if (viewController == self) {
+                break;
+            }
+            previous = viewController;
+        }
+        [self.navigationController stp_popToViewController:previous animated:YES completion:completion];
+    }
+}
+
 @end
 
 @implementation STPPaymentMethodsViewController (Private)
diff --git a/Stripe/STPShippingAddressViewController.m b/Stripe/STPShippingAddressViewController.m
index 870f742055..937bf36e74 100644
--- a/Stripe/STPShippingAddressViewController.m
+++ b/Stripe/STPShippingAddressViewController.m
@@ -263,8 +263,20 @@
     [self presentViewController:alertController animated:YES completion:nil];
 }
 
-- (void)dismissWithHostViewController:(UIViewController *)hostViewController {
-    [self.navigationController stp_popToViewController:hostViewController animated:YES completion:nil];
+- (void)dismissWithCompletion:(STPVoidBlock)completion {
+    if ([self stp_isAtRootOfNavigationController]) {
+        [self.presentingViewController dismissViewControllerAnimated:YES completion:completion];
+    }
+    else {
+        UIViewController *previous = self.navigationController.viewControllers.firstObject;
+        for (UIViewController *viewController in self.navigationController.viewControllers) {
+            if (viewController == self) {
+                break;
+            }
+            previous = viewController;
+        }
+        [self.navigationController stp_popToViewController:previous animated:YES completion:completion];
+    }
 }
 
 #pragma mark - STPAddressViewModelDelegate