From 9608865c852e8d50709f9842d29b08994081cb64 Mon Sep 17 00:00:00 2001
From: Yuwnei Xiao <xiaoluoboding@gmail.com>
Date: Fri, 3 Nov 2023 11:45:25 +0800
Subject: [PATCH] feat: make promise return promise

---
 packages/state.ts | 66 ++++++++++++++++++++++++++++++++++++++---------
 packages/types.ts | 11 ++++----
 2 files changed, 59 insertions(+), 18 deletions(-)

diff --git a/packages/state.ts b/packages/state.ts
index bb06722..acbcb4d 100644
--- a/packages/state.ts
+++ b/packages/state.ts
@@ -92,21 +92,15 @@ class Observer {
   }
 
   message = (message: string | Component, data?: ExternalToast) => {
-    const id = data?.id || toastsCounter++
-    this.publish({ ...data, id, title: message })
-    return id
+    return this.create({ ...data, message })
   }
 
   error = (message: string | Component, data?: ExternalToast) => {
-    const id = data?.id || toastsCounter++
-    this.publish({ ...data, id, type: 'error', title: message })
-    return id
+    return this.create({ ...data, type: 'error', message })
   }
 
   success = (message: string | Component, data?: ExternalToast) => {
-    const id = data?.id || toastsCounter++
-    this.publish({ ...data, id, type: 'success', title: message })
-    return id
+    return this.create({ ...data, type: 'success', message })
   }
 
   info = (message: string | Component, data?: ExternalToast) => {
@@ -121,9 +115,57 @@ class Observer {
     return this.create({ ...data, type: 'loading', message })
   }
 
-  promise = (promise: PromiseT, data?: PromiseData) => {
-    const id = data?.id || toastsCounter++
-    this.publish({ ...data, promise, id })
+  promise = <ToastData>(
+    promise: PromiseT<ToastData>,
+    data?: PromiseData<ToastData>
+  ) => {
+    if (!data) {
+      // Nothing to show
+      return
+    }
+
+    let id: string | number | undefined = undefined
+    if (data.loading !== undefined) {
+      id = this.create({
+        ...data,
+        promise,
+        type: 'loading',
+        message: data.loading
+      })
+    }
+
+    const p = promise instanceof Promise ? promise : promise()
+
+    let shouldDismiss = id !== undefined
+
+    p.then((promiseData) => {
+      if (data.success !== undefined) {
+        shouldDismiss = false
+        const message =
+          typeof data.success === 'function'
+            ? data.success(promiseData)
+            : data.success
+        this.create({ id, type: 'success', message })
+      }
+    })
+      .catch((error) => {
+        if (data.error !== undefined) {
+          shouldDismiss = false
+          const message =
+            typeof data.error === 'function' ? data.error(error) : data.error
+          this.create({ id, type: 'error', message })
+        }
+      })
+      .finally(() => {
+        if (shouldDismiss) {
+          // Toast is still in load state (and will be indefinitely — dismiss it)
+          this.dismiss(id)
+          id = undefined
+        }
+
+        data.finally?.()
+      })
+
     return id
   }
 
diff --git a/packages/types.ts b/packages/types.ts
index cdbde2c..ac84227 100644
--- a/packages/types.ts
+++ b/packages/types.ts
@@ -11,12 +11,11 @@ export type ToastTypes =
 
 export type PromiseT<Data = any> = Promise<Data> | (() => Promise<Data>)
 
-export type PromiseData = ExternalToast & {
-  loading: string | Component
-  // success: string | Component | ((data: any) => Component | string)
-  // error: string | Component | ((error: any) => Component | string)
-  success: (data: any) => Component | string
-  error: (error: any) => Component | string
+export type PromiseData<ToastData = any> = ExternalToast & {
+  loading?: string | Component
+  success?: (data: ToastData) => string | Component
+  error?: (data: ToastData) => string | Component
+  finally?: () => void | Promise<void>
 }
 
 export interface ToastT {