diff --git a/interfaces.d.ts b/interfaces.d.ts
index 4e1beaf..f668fb4 100644
--- a/interfaces.d.ts
+++ b/interfaces.d.ts
@@ -8,8 +8,11 @@ export interface ReactFocusLockProps<ChildrenType = React.ReactNode, LockProps=R
    * if true, will return focus to the previous position on trap disable.
    * Optionally, can pass focus options instead of `true` to control the focus
    * more precisely (ie. `{ preventScroll: true }`)
+   *
+   * can also accept a function with the first argument equals to element focus will be returned to
+   * in order to
    */
-  returnFocus?: boolean | FocusOptions;
+  returnFocus?: boolean | FocusOptions | ((returnTo: Element)=> boolean | FocusOptions);
 
   /**
    * @deprecated Use persistentFocus=false instead
diff --git a/src/Lock.js b/src/Lock.js
index 4ccbfcf..0319ec8 100644
--- a/src/Lock.js
+++ b/src/Lock.js
@@ -3,10 +3,10 @@ import {
   node, bool, string, any, arrayOf, oneOfType, object, func,
 } from 'prop-types';
 import * as constants from 'focus-lock/constants';
-import { useMergeRefs } from 'use-callback-ref';
+import {useMergeRefs} from 'use-callback-ref';
 
-import { hiddenGuard } from './FocusGuard';
-import { mediumFocus, mediumBlur, mediumSidecar } from './medium';
+import {hiddenGuard} from './FocusGuard';
+import {mediumFocus, mediumBlur, mediumSidecar} from './medium';
 
 const emptyArray = [];
 
@@ -60,17 +60,20 @@ const FocusLock = React.forwardRef(function FocusLockUI(props, parentRef) {
   }, [onDeactivationCallback]);
 
   const returnFocus = React.useCallback((allowDefer) => {
-    const { current } = originalFocusedElement;
-    if (Boolean(shouldReturnFocus) && current && current.focus) {
-      const focusOptions = typeof shouldReturnFocus === 'object' ? shouldReturnFocus : undefined;
-      originalFocusedElement.current = null;
-
-      if (allowDefer) {
-        // React might return focus after update
-        // it's safer to defer the action
-        Promise.resolve().then(() => current.focus(focusOptions));
-      } else {
-        current.focus(focusOptions);
+    const {current: returnFocusTo} = originalFocusedElement;
+    if (returnFocusTo && returnFocusTo.focus) {
+      const howToReturnFocus = typeof shouldReturnFocus === 'function' ? shouldReturnFocus(returnFocusTo) : shouldReturnFocus;
+      if (Boolean(howToReturnFocus)) {
+        const focusOptions = typeof howToReturnFocus === 'object' ? howToReturnFocus : undefined;
+        originalFocusedElement.current = null;
+
+        if (allowDefer) {
+          // React might return focus after update
+          // it's safer to defer the action
+          Promise.resolve().then(() => returnFocusTo.focus(focusOptions));
+        } else {
+          returnFocusTo.focus(focusOptions);
+        }
       }
     }
   }, [shouldReturnFocus]);
@@ -123,8 +126,8 @@ const FocusLock = React.forwardRef(function FocusLockUI(props, parentRef) {
   return (
     <React.Fragment>
       {hasLeadingGuards && [
-        <div key="guard-first" data-focus-guard tabIndex={disabled ? -1 : 0} style={hiddenGuard} />, // nearest focus guard
-        <div key="guard-nearest" data-focus-guard tabIndex={disabled ? -1 : 1} style={hiddenGuard} />, // first tabbed element guard
+        <div key="guard-first" data-focus-guard tabIndex={disabled ? -1 : 0} style={hiddenGuard}/>, // nearest focus guard
+        <div key="guard-nearest" data-focus-guard tabIndex={disabled ? -1 : 1} style={hiddenGuard}/>, // first tabbed element guard
       ]}
       {!disabled && (
         <SideCar
@@ -153,7 +156,7 @@ const FocusLock = React.forwardRef(function FocusLockUI(props, parentRef) {
       </Container>
       {
         hasTailingGuards
-        && <div data-focus-guard tabIndex={disabled ? -1 : 0} style={hiddenGuard} />
+        && <div data-focus-guard tabIndex={disabled ? -1 : 0} style={hiddenGuard}/>
       }
     </React.Fragment>
   );