chore(storybook): standardize down-state decorator

Author: castastrophe

.storybook/decorators/down-state.js

@@ -1,19 +1,40 @@
-export const withDownStateDimensionCapture = (selector) => (Story, context) => {
-  const captureDownStateDimensions = () => {
-    const components = document.querySelectorAll(selector);
-    components.forEach((component) => {
-      const { width, height } = component.getBoundingClientRect();
-      component.style.setProperty('--spectrum-downstate-width', `${width}px`);
-      component.style.setProperty('--spectrum-downstate-height', `${height}px`);
-    });
-  };
+import { makeDecorator, useEffect } from "@storybook/preview-api";
+import { fetchContainers } from "./helpers.js";
 
-  document.addEventListener("DOMContentLoaded", () => {
-    // Wait to make sure the story is fully rendered (otherwise width/height can be wrong)
-    setTimeout(() => {
-        captureDownStateDimensions();
-    }, 100);
-  });
+export const withDownStateDimensionCapture = makeDecorator({
+  name: "withDownStateDimensionCapture",
+  parameterName: "downState",
+  wrapper: (StoryFn, context) => {
+    const { args = {}, parameters = {}, viewMode, id } = context;
+    
+    /* Selectors are defined in the downState parameter */
+    const {
+      // Fall back to the rootClass if no selectors are provided
+      selectors = args.rootClass ? [ args.rootClass ] : []
+    } = parameters.downState ?? {};
 
-  return Story(context);
-};
+    /**
+     * This effect will run after the component is rendered and will capture the dimensions of the
+     * components that are specified in the selectors array. It will then set the custom properties
+     * --spectrum-downstate-width and --spectrum-downstate-height on the component to the width and
+     * height of the component respectively. This will allow the downstate to be calculated correctly
+     * in the CSS.
+     */
+    useEffect(() => {
+      if (selectors.length === 0) return;
+
+      for (const container of fetchContainers(id, viewMode === "docs")) {
+        for (const selector of selectors) {
+          const components = [...container.querySelectorAll(selector)];
+          for (const component of components) {
+            const { width, height } = component.getBoundingClientRect();
+            component.style.setProperty("--spectrum-downstate-width", `${width}px`);
+            component.style.setProperty("--spectrum-downstate-height", `${height}px`);
+          }
+        }
+      }
+    }, [selectors]);
+
+    return StoryFn(context);
+  },
+});

components/button/stories/button.stories.js

@@ -89,11 +89,17 @@ export default {
 			type: "figma",
 			url: "https://www.figma.com/design/Mngz9H7WZLbrCvGQf3GnsY/S2-%2F-Desktop?node-id=707-2774",
 		},
+		html: {
+			root: "#render-root"
+		},
+		downState: {
+			selectors: [".spectrum-Button:not(:disabled)"],
+		}
 		packageJson,
 		metadata,
 	},
 	decorators: [
-		withDownStateDimensionCapture('.spectrum-Button:not(:disabled)')
+		withDownStateDimensionCapture,
 	],
 	tags: ["!autodocs"],
 };

components/picker/stories/picker.stories.js

@@ -142,9 +142,12 @@ export default {
 		},
 		packageJson,
 		metadata,
+		downState: {
+			selectors: [".spectrum-Picker:not(:disabled, .is-disabled, .is-loading)"],
+		},
 	},
 	decorators: [
-		withDownStateDimensionCapture(".spectrum-Picker:not(:disabled, .is-disabled, .is-loading)"),
+		withDownStateDimensionCapture,
 	],
 };