chore: generalise setStyle with cleanup and fix "Cannot set property focus..."

packages/react-aria/src/overlays/usePreventScroll.ts

@@ -11,12 +11,12 @@
  */
 
 import {chain} from '../utils/chain';
-
 import {getActiveElement, getEventTarget} from '../utils/shadowdom/DOMFunctions';
 import {getNonce} from '../utils/getNonce';
 import {getScrollParent} from '../utils/getScrollParent';
 import {isIOS, isWebKit} from '../utils/platform';
 import {isScrollable} from '../utils/isScrollable';
+import {setStyle} from '../utils/domHelpers';
 import {useLayoutEffect} from '../utils/useLayoutEffect';
 import {willOpenKeyboard} from '../utils/keyboard';
 
@@ -70,8 +70,8 @@ function preventScrollStandard() {
     scrollbarWidth > 0 &&
       // Use scrollbar-gutter when supported because it also works for fixed positioned elements.
       ('scrollbarGutter' in document.documentElement.style
-        ? setStyle(document.documentElement, 'scrollbarGutter', 'stable')
-        : setStyle(document.documentElement, 'paddingRight', `${scrollbarWidth}px`)),
+        ? setStyle(document.documentElement, 'scrollbar-gutter', 'stable')
+        : setStyle(document.documentElement, 'padding-right', `${scrollbarWidth}px`)),
     setStyle(document.documentElement, 'overflow', 'hidden')
   );
 }
@@ -194,18 +194,22 @@ function preventScrollMobileWebKit() {
 
   // Override programmatic focus to scroll into view without scrolling the whole page.
   let focus = HTMLElement.prototype.focus;
-  HTMLElement.prototype.focus = function (opts) {
-    // Track whether the keyboard was already visible before.
-    let activeElement = getActiveElement();
-    let wasKeyboardVisible = activeElement != null && willOpenKeyboard(activeElement);
-
-    // Focus the element without scrolling the page.
-    focus.call(this, {...opts, preventScroll: true});
-
-    if (!opts || !opts.preventScroll) {
-      scrollIntoViewWhenReady(this, wasKeyboardVisible);
+  Reflect.defineProperty(HTMLElement.prototype, 'focus', {
+    configurable: true,
+    writable: true,
+    value: function (opts?: FocusOptions) {
+      // Track whether the keyboard was already visible before.
+      let activeElement = getActiveElement();
+      let wasKeyboardVisible = activeElement != null && willOpenKeyboard(activeElement);
+
+      // Focus the element without scrolling the page.
+      focus.call(this, {...opts, preventScroll: true});
+
+      if (!opts || !opts.preventScroll) {
+        scrollIntoViewWhenReady(this, wasKeyboardVisible);
+      }
     }
-  };
+  });
 
   let removeEvents = chain(
     addEvent(document, 'touchstart', onTouchStart, {passive: false, capture: true}),
@@ -217,17 +221,11 @@ function preventScrollMobileWebKit() {
     restoreOverflow();
     removeEvents();
     style.remove();
-    HTMLElement.prototype.focus = focus;
-  };
-}
-
-// Sets a CSS property on an element, and returns a function to revert it to the previous value.
-function setStyle(element: HTMLElement, style: string, value: string) {
-  let cur = element.style[style];
-  element.style[style] = value;
-
-  return () => {
-    element.style[style] = cur;
+    Reflect.defineProperty(HTMLElement.prototype, 'focus', {
+      configurable: true,
+      writable: true,
+      value: focus
+    });
   };
 }
 

packages/react-aria/src/utils/domHelpers.ts

@@ -91,3 +91,41 @@ export function addEvent<T extends EventTarget, K extends keyof EventMapType<Exc
     }
   };
 }
+
+/**
+ * Sets a CSS property on an element and returns a cleanup function.
+ */
+export function setStyle(
+  target: HTMLElement | HTMLElement[] | null,
+  property: string,
+  value: string,
+  priority?: string
+): () => void {
+  if (target == null) {
+    return () => {};
+  }
+
+  let restore = new Array<Function>();
+  let styleTargets = Array.isArray(target) ? target : [target];
+
+  for (let styleTarget of styleTargets) {
+    let initialValue = styleTarget.style.getPropertyValue(property);
+    let initialPriority = styleTarget.style.getPropertyPriority(property);
+
+    styleTarget.style.setProperty(property, value, priority);
+
+    restore.unshift(() => {
+      if (initialValue) {
+        styleTarget.style.setProperty(property, initialValue, initialPriority);
+      } else {
+        styleTarget.style.removeProperty(property);
+      }
+    });
+  }
+
+  return () => {
+    for (let cleanup of restore) {
+      cleanup();
+    }
+  };
+}

packages/react-aria/test/utils/domHelpers.test.js

@@ -13,7 +13,7 @@
 import {act} from 'react-dom/test-utils';
 import {enableShadowDOM} from 'react-stately/private/flags/flags';
 import {getActiveElement} from '../../src/utils/shadowdom/DOMFunctions';
-import {getOwnerDocument, getOwnerWindow} from '../../src/utils/domHelpers';
+import {getOwnerDocument, getOwnerWindow, setStyle} from '../../src/utils/domHelpers';
 
 describe('getOwnerDocument', () => {
   beforeAll(() => {
@@ -191,3 +191,59 @@ describe('getActiveElement', () => {
     iframe.remove();
   });
 });
+
+describe('setStyle', () => {
+  it('returns a no-op cleanup and does not throw when the target is null', () => {
+    const cleanup = setStyle(null, 'opacity', '0');
+    expect(cleanup).toBeInstanceOf(Function);
+    expect(() => cleanup()).not.toThrow();
+  });
+
+  it('sets a CSS property on a single element and removes the property on cleanup when there was no initial value', () => {
+    const el = document.createElement('div');
+    const cleanup = setStyle(el, 'opacity', '0');
+    expect(el.style.getPropertyValue('opacity')).toBe('0');
+
+    cleanup();
+    expect(el.style.getPropertyValue('opacity')).toBe('');
+    expect(el.getAttribute('style')).toBeFalsy();
+  });
+
+  it('restores the previous value on cleanup when one existed', () => {
+    const el = document.createElement('div');
+    el.style.setProperty('opacity', '0.5');
+
+    const cleanup = setStyle(el, 'opacity', '0');
+    expect(el.style.getPropertyValue('opacity')).toBe('0');
+
+    cleanup();
+    expect(el.style.getPropertyValue('opacity')).toBe('0.5');
+  });
+
+  it('applies the given priority and restores the previous priority on cleanup', () => {
+    const el = document.createElement('div');
+    el.style.setProperty('color', 'red', 'important');
+
+    const cleanup = setStyle(el, 'color', 'blue', 'important');
+    expect(el.style.getPropertyValue('color')).toBe('blue');
+    expect(el.style.getPropertyPriority('color')).toBe('important');
+
+    cleanup();
+    expect(el.style.getPropertyValue('color')).toBe('red');
+    expect(el.style.getPropertyPriority('color')).toBe('important');
+  });
+
+  it('ets the property on every element in an array and restores on cleanup, preserving each prior value', () => {
+    const withPrior = document.createElement('div');
+    withPrior.style.setProperty('display', 'flex');
+    const withoutPrior = document.createElement('div');
+
+    const cleanup = setStyle([withPrior, withoutPrior], 'display', 'none');
+    expect(withPrior.style.getPropertyValue('display')).toBe('none');
+    expect(withoutPrior.style.getPropertyValue('display')).toBe('none');
+
+    cleanup();
+    expect(withPrior.style.getPropertyValue('display')).toBe('flex');
+    expect(withoutPrior.style.getPropertyValue('display')).toBe('');
+  });
+});