From c4a14a4fb370c7628e4cc5861e31cc64a66b64b0 Mon Sep 17 00:00:00 2001
From: Stefan <stefan@gigmade.com>
Date: Wed, 8 Jan 2020 09:13:18 +0100
Subject: [PATCH] feat(useBeforeUnload): allow passing a dirty function (#842)

---
 docs/useBeforeUnload.md           | 27 +++++++++++++++++++++++++++
 src/useBeforeUnload.ts            | 27 ++++++++++++++++++---------
 stories/useBeforeUnload.story.tsx | 22 +++++++++++++++++++---
 3 files changed, 64 insertions(+), 12 deletions(-)

diff --git a/docs/useBeforeUnload.md b/docs/useBeforeUnload.md
index f95cfe9141..ca7dbe4a0f 100644
--- a/docs/useBeforeUnload.md
+++ b/docs/useBeforeUnload.md
@@ -5,6 +5,8 @@ React side-effect hook that shows browser alert when user try to reload or close
 
 ## Usage
 
+### Boolean check
+
 ```jsx
 import {useBeforeUnload} from 'react-use';
 
@@ -20,3 +22,28 @@ const Demo = () => {
   );
 };
 ```
+
+### Function check
+
+Note: Since every `dirtyFn` change registers a new callback, you should use
+[refs](https://reactjs.org/docs/hooks-faq.html#how-to-read-an-often-changing-value-from-usecallback)
+if your test value changes often.
+
+```jsx
+import {useBeforeUnload} from 'react-use';
+
+const Demo = () => {
+  const [dirty, toggleDirty] = useToggle(false);
+  const dirtyFn = useCallback(() => {
+    return dirty;
+  }, [dirty]);
+  useBeforeUnload(dirtyFn, 'You have unsaved changes, are you sure?');
+
+  return (
+    <div>
+      {dirty && <p>Try to reload or close tab</p>}
+      <button onClick={() => toggleDirty()}>{dirty ? 'Disable' : 'Enable'}</button>
+    </div>
+  );
+};
+```
diff --git a/src/useBeforeUnload.ts b/src/useBeforeUnload.ts
index f209775959..83d9410f5d 100644
--- a/src/useBeforeUnload.ts
+++ b/src/useBeforeUnload.ts
@@ -1,12 +1,14 @@
-import { useEffect } from 'react';
+import { useCallback, useEffect } from 'react';
 
-const useBeforeUnload = (enabled: boolean = true, message?: string) => {
-  useEffect(() => {
-    if (!enabled) {
-      return;
-    }
+const useBeforeUnload = (enabled: boolean | (() => boolean) = true, message?: string) => {
+  const handler = useCallback(
+    (event: BeforeUnloadEvent) => {
+      const finalEnabled = typeof enabled === 'function' ? enabled() : true;
+
+      if (!finalEnabled) {
+        return;
+      }
 
-    const handler = (event: BeforeUnloadEvent) => {
       event.preventDefault();
 
       if (message) {
@@ -14,12 +16,19 @@ const useBeforeUnload = (enabled: boolean = true, message?: string) => {
       }
 
       return message;
-    };
+    },
+    [enabled, message]
+  );
+
+  useEffect(() => {
+    if (!enabled) {
+      return;
+    }
 
     window.addEventListener('beforeunload', handler);
 
     return () => window.removeEventListener('beforeunload', handler);
-  }, [message, enabled]);
+  }, [enabled, handler]);
 };
 
 export default useBeforeUnload;
diff --git a/stories/useBeforeUnload.story.tsx b/stories/useBeforeUnload.story.tsx
index d87a14bd96..1fc25008f4 100644
--- a/stories/useBeforeUnload.story.tsx
+++ b/stories/useBeforeUnload.story.tsx
@@ -1,9 +1,9 @@
 import { storiesOf } from '@storybook/react';
-import * as React from 'react';
+import React, { useCallback } from 'react';
 import { useBeforeUnload, useToggle } from '../src';
 import ShowDocs from './util/ShowDocs';
 
-const Demo = () => {
+const DemoBool = () => {
   const [dirty, toggleDirty] = useToggle(false);
   useBeforeUnload(dirty, 'You have unsaved changes, are you sure?');
 
@@ -15,6 +15,22 @@ const Demo = () => {
   );
 };
 
+const DemoFunc = () => {
+  const [dirty, toggleDirty] = useToggle(false);
+  const dirtyFn = useCallback(() => {
+    return dirty;
+  }, [dirty]);
+  useBeforeUnload(dirtyFn, 'You have unsaved changes, are you sure?');
+
+  return (
+    <div>
+      {dirty && <p>Try to reload or close tab</p>}
+      <button onClick={() => toggleDirty()}>{dirty ? 'Disable' : 'Enable'}</button>
+    </div>
+  );
+};
+
 storiesOf('Side effects|useBeforeUnload', module)
   .add('Docs', () => <ShowDocs md={require('../docs/useBeforeUnload.md')} />)
-  .add('Demo', () => <Demo />);
+  .add('Demo (boolean)', () => <DemoBool />)
+  .add('Demo (function)', () => <DemoFunc />);