allow passing a target document through contexts

Author: gabriellsh

packages/css-in-js/src/sheet.ts

@@ -1,20 +1,27 @@
 import hash from '@emotion/hash';
 
 const elementId = 'rcx-styles';
-let element: HTMLStyleElement;
-const getStyleTag = (): HTMLStyleElement => {
+const documentStyleElementMap = new Map<Document, HTMLStyleElement>();
+
+const getStyleTag = (document: Document): HTMLStyleElement => {
+  let element = documentStyleElementMap.get(document);
+  console.trace('getStyleTag', document, element);
   if (!element) {
     const el = document.getElementById(elementId) as HTMLStyleElement;
     if (el) {
-      element = el;
-      return element;
+      documentStyleElementMap.set(document, el);
+      return el;
     }
   }
 
   if (!element) {
     element = document.createElement('style');
+
+    documentStyleElementMap.set(document, element);
+
     element.id = elementId;
     element.appendChild(document.createTextNode(''));
+
     if (document.head) {
       document.head.appendChild(element);
     }
@@ -23,10 +30,13 @@ const getStyleTag = (): HTMLStyleElement => {
   return element;
 };
 
-let styleSheet: CSSStyleSheet;
-const getStyleSheet = (): CSSStyleSheet => {
+// let styleSheet: CSSStyleSheet;
+const documentStyleSheetMap = new Map<Document, CSSStyleSheet>();
+const getStyleSheet = (document: Document): CSSStyleSheet => {
+  let styleSheet = documentStyleSheetMap.get(document);
+
   if (!styleSheet) {
-    const styleTag = getStyleTag();
+    const styleTag = getStyleTag(document);
     const _styleSheet =
       styleTag.sheet ||
       Array.from(document.styleSheets).find(
@@ -40,24 +50,29 @@ const getStyleSheet = (): CSSStyleSheet => {
     styleSheet = _styleSheet;
   }
 
+  documentStyleSheetMap.set(document, styleSheet);
+
   return styleSheet;
 };
 
-type RuleAttacher = (rules: string) => () => void;
+type RuleAttacher = (
+  rules: string,
+  options: { document: Document },
+) => () => void;
 
 const discardRules: RuleAttacher = () => () => undefined;
 
-const attachRulesIntoElement: RuleAttacher = (rules) => {
-  const element = getStyleTag();
+const attachRulesIntoElement: RuleAttacher = (rules, { document }) => {
+  const element = getStyleTag(document);
 
   const textNode = document.createTextNode(rules);
   element.appendChild(textNode);
 
   return () => textNode.remove();
 };
 
-const attachRulesIntoStyleSheet: RuleAttacher = (rules) => {
-  const styleSheet = getStyleSheet();
+const attachRulesIntoStyleSheet: RuleAttacher = (rules, { document }) => {
+  const styleSheet = getStyleSheet(document);
   const index = styleSheet.insertRule(
     `@media all{${rules}}`,
     styleSheet.cssRules.length,
@@ -88,11 +103,14 @@ const wrapReferenceCounting = (attacher: RuleAttacher): RuleAttacher => {
     window.queueMicrotask(fn);
   };
 
-  const enhancedAttacher: RuleAttacher = (content: string) => {
+  const enhancedAttacher: RuleAttacher = (content, options) => {
     const id = hash(content);
 
     if (!refs[id]) {
-      const detach = attacher(content);
+      const detach = attacher(content, {
+        ...options,
+        document: options.document ? options.document : window.document,
+      });
       let count = 0;
 
       const ref = (): void => {

packages/fuselage/src/hooks/useStyle.ts

@@ -5,7 +5,21 @@ import {
   transpile,
   attachRules,
 } from '@rocket.chat/css-in-js';
-import { useDebugValue, useInsertionEffect, useMemo } from 'react';
+import {
+  createContext,
+  useContext,
+  useDebugValue,
+  useInsertionEffect,
+  useMemo,
+} from 'react';
+
+export const StyleOptions = createContext<{
+  document?: Document;
+}>({ document: window.document });
+
+export const useStyleOptions = () => {
+  return useContext(StyleOptions);
+};
 
 export const useStyle = (cssFn: cssFn | undefined, arg: unknown) => {
   const content = useMemo(() => (cssFn ? cssFn(arg) : undefined), [arg, cssFn]);
@@ -20,19 +34,21 @@ export const useStyle = (cssFn: cssFn | undefined, arg: unknown) => {
 
   useDebugValue(className);
 
+  const { document } = useStyleOptions();
+
   useInsertionEffect(() => {
     if (!content || !className) {
       return;
     }
 
     const escapedClassName = escapeName(className);
     const transpiledContent = transpile(`.${escapedClassName}`, content);
-    const detach = attachRules(transpiledContent);
+    const detach = attachRules(transpiledContent, { document });
 
     return () => {
       setTimeout(detach, 1000);
     };
-  }, [className, content]);
+  }, [className, content, document]);
 
   return className;
 };

packages/fuselage/src/index.ts

@@ -3,5 +3,7 @@ import './index.scss';
 export * from './components';
 export * from './styleTokens';
 
+export { StyleOptions } from './hooks/useStyle';
+
 export { Palette, __setThrowErrorOnInvalidToken__ } from './Theme';
 export { useArrayLikeClassNameProp } from './hooks/useArrayLikeClassNameProp';

packages/styled/src/index.ts

@@ -1 +1 @@
-export { default } from './styled';
+export { default, StyledOptions } from './styled';

packages/styled/src/styled.ts

@@ -20,8 +20,18 @@ import {
   Fragment,
   useDebugValue,
   useInsertionEffect,
+  createContext,
+  useContext,
 } from 'react';
 
+export const StyledOptions = createContext<{
+  document?: Document;
+}>({ document: window.document });
+
+const useStyledOptions = () => {
+  return useContext(StyledOptions);
+};
+
 export const attachClassName = <P extends { className?: string }>(
   props: P,
   additionalClassName: string,
@@ -85,18 +95,19 @@ const styled =
 
           useDebugValue(computedClassName);
 
+          const { document } = useStyledOptions();
           useInsertionEffect(() => {
             const escapedClassName = escapeName(computedClassName);
             const transpiledContent = transpile(
               `.${escapedClassName}`,
               content,
             );
-            const detach = attachRules(transpiledContent);
+            const detach = attachRules(transpiledContent, { document });
 
             return () => {
               setTimeout(detach, 1000);
             };
-          }, [computedClassName, content]);
+          }, [computedClassName, content, document]);
 
           const newProps = attachClassName(
             { ref, ...props },