refactor: consolidate Theme-related code

packages/styled-components/src/base.ts

@@ -6,7 +6,6 @@ import keyframes from './constructors/keyframes';
 /* Import Higher Order Components */
 import withTheme from './hoc/withTheme';
 /* Import hooks */
-import useTheme from './hooks/useTheme';
 import ServerStyleSheet from './models/ServerStyleSheet';
 import {
   IStyleSheetContext,
@@ -17,7 +16,7 @@ import {
   StyleSheetManager,
 } from './models/StyleSheetManager';
 /* Import components */
-import ThemeProvider, { ThemeConsumer, ThemeContext } from './models/ThemeProvider';
+import ThemeProvider, { ThemeConsumer, ThemeContext, useTheme } from './models/ThemeProvider';
 import isStyledComponent from './utils/isStyledComponent';
 
 /* Warning if you've imported this file on React Native */

packages/styled-components/src/constructors/createGlobalStyle.ts

@@ -1,8 +1,8 @@
 import React from 'react';
 import { STATIC_EXECUTION_CONTEXT } from '../constants';
 import GlobalStyle from '../models/GlobalStyle';
-import { useStyleSheet, useStylis } from '../models/StyleSheetManager';
-import { DefaultTheme, ThemeContext } from '../models/ThemeProvider';
+import { useStyleSheetContext } from '../models/StyleSheetManager';
+import { DefaultTheme, useTheme } from '../models/ThemeProvider';
 import StyleSheet from '../sheet';
 import { ExecutionContext, ExecutionProps, Interpolation, Stringifier, Styles } from '../types';
 import { checkDynamicCreation } from '../utils/checkDynamicCreation';
@@ -23,10 +23,9 @@ export default function createGlobalStyle<Props extends object>(
   }
 
   const GlobalStyleComponent: React.ComponentType<ExecutionProps & Props> = props => {
-    const styleSheet = useStyleSheet();
-    const stylis = useStylis();
-    const theme = React.useContext(ThemeContext);
-    const instanceRef = React.useRef(styleSheet.allocateGSInstance(styledComponentId));
+    const ssc = useStyleSheetContext();
+    const theme = useTheme();
+    const instanceRef = React.useRef(ssc.styleSheet.allocateGSInstance(styledComponentId));
 
     const instance = instanceRef.current;
 
@@ -45,18 +44,18 @@ export default function createGlobalStyle<Props extends object>(
       );
     }
 
-    if (styleSheet.server) {
-      renderStyles(instance, props, styleSheet, theme, stylis);
+    if (ssc.styleSheet.server) {
+      renderStyles(instance, props, ssc.styleSheet, theme, ssc.stylis);
     }
 
     if (!__SERVER__) {
       // @ts-expect-error still using React 17 types for the time being
       (React.useInsertionEffect || React.useLayoutEffect)(() => {
-        if (!styleSheet.server) {
-          renderStyles(instance, props, styleSheet, theme, stylis);
-          return () => globalStyle.removeStyles(instance, styleSheet);
+        if (!ssc.styleSheet.server) {
+          renderStyles(instance, props, ssc.styleSheet, theme, ssc.stylis);
+          return () => globalStyle.removeStyles(instance, ssc.styleSheet);
         }
-      }, [instance, props, styleSheet, theme, stylis]);
+      }, [instance, props, ssc.styleSheet, theme, ssc.stylis]);
     }
 
     return null;

packages/styled-components/src/hoc/withTheme.tsx

@@ -1,5 +1,5 @@
 import React from 'react';
-import { ThemeContext } from '../models/ThemeProvider';
+import { useTheme } from '../models/ThemeProvider';
 import { AnyComponent, ExecutionProps } from '../types';
 import determineTheme from '../utils/determineTheme';
 import getComponentName from '../utils/getComponentName';
@@ -8,7 +8,7 @@ import hoist from '../utils/hoist';
 export default function withTheme<T extends AnyComponent>(Component: T) {
   const WithTheme = React.forwardRef<T, JSX.LibraryManagedAttributes<T, ExecutionProps>>(
     (props, ref) => {
-      const theme = React.useContext(ThemeContext);
+      const theme = useTheme();
       const themeProp = determineTheme(props, theme, Component.defaultProps);
 
       if (process.env.NODE_ENV !== 'production' && themeProp === undefined) {

packages/styled-components/src/models/StyleSheetManager.tsx

@@ -4,14 +4,19 @@ import StyleSheet from '../sheet';
 import { ShouldForwardProp, Stringifier } from '../types';
 import createStylisInstance from '../utils/stylis';
 
+export const mainSheet: StyleSheet = new StyleSheet();
+export const mainStylis: Stringifier = createStylisInstance();
+
 export type IStyleSheetContext = {
   shouldForwardProp?: ShouldForwardProp<'web'>;
-  styleSheet?: StyleSheet;
+  styleSheet: StyleSheet;
+  stylis: Stringifier;
 };
 
 export const StyleSheetContext = React.createContext<IStyleSheetContext>({
   shouldForwardProp: undefined,
-  styleSheet: undefined,
+  styleSheet: mainSheet,
+  stylis: mainStylis,
 });
 
 export const StyleSheetConsumer = StyleSheetContext.Consumer;
@@ -20,19 +25,8 @@ export type IStylisContext = Stringifier | void;
 export const StylisContext = React.createContext<IStylisContext>(undefined);
 export const StylisConsumer = StylisContext.Consumer;
 
-export const mainSheet: StyleSheet = new StyleSheet();
-export const mainStylis: Stringifier = createStylisInstance();
-
-export function useShouldForwardProp() {
-  return useContext(StyleSheetContext).shouldForwardProp;
-}
-
-export function useStyleSheet(): StyleSheet {
-  return useContext(StyleSheetContext).styleSheet || mainSheet;
-}
-
-export function useStylis(): Stringifier {
-  return useContext(StylisContext) || mainStylis;
+export function useStyleSheetContext() {
+  return useContext(StyleSheetContext);
 }
 
 export type IStyleSheetManager = React.PropsWithChildren<{
@@ -85,10 +79,10 @@ export type IStyleSheetManager = React.PropsWithChildren<{
 
 export function StyleSheetManager(props: IStyleSheetManager): JSX.Element {
   const [plugins, setPlugins] = useState(props.stylisPlugins);
-  const contextStyleSheet = useStyleSheet();
+  const { styleSheet } = useStyleSheetContext();
 
-  const styleSheet = useMemo(() => {
-    let sheet = contextStyleSheet;
+  const resolvedStyleSheet = useMemo(() => {
+    let sheet = styleSheet;
 
     if (props.sheet) {
       sheet = props.sheet;
@@ -101,7 +95,7 @@ export function StyleSheetManager(props: IStyleSheetManager): JSX.Element {
     }
 
     return sheet;
-  }, [props.disableCSSOMInjection, props.sheet, props.target]);
+  }, [props.disableCSSOMInjection, props.sheet, props.target, styleSheet]);
 
   const stylis = useMemo(
     () =>
@@ -117,7 +111,9 @@ export function StyleSheetManager(props: IStyleSheetManager): JSX.Element {
   }, [props.stylisPlugins]);
 
   return (
-    <StyleSheetContext.Provider value={{ shouldForwardProp: props.shouldForwardProp, styleSheet }}>
+    <StyleSheetContext.Provider
+      value={{ shouldForwardProp: props.shouldForwardProp, styleSheet: resolvedStyleSheet, stylis }}
+    >
       <StylisContext.Provider value={stylis}>{props.children}</StylisContext.Provider>
     </StyleSheetContext.Provider>
   );

packages/styled-components/src/models/StyledComponent.ts

@@ -1,4 +1,4 @@
-import React, { createElement, Ref, useContext, useDebugValue } from 'react';
+import React, { createElement, Ref, useDebugValue } from 'react';
 import { SC_VERSION } from '../constants';
 import type {
   AnyComponent,
@@ -29,8 +29,8 @@ import isTag from '../utils/isTag';
 import joinStrings from '../utils/joinStrings';
 import merge from '../utils/mixinDeep';
 import ComponentStyle from './ComponentStyle';
-import { useShouldForwardProp, useStyleSheet, useStylis } from './StyleSheetManager';
-import { DefaultTheme, ThemeContext } from './ThemeProvider';
+import { useStyleSheetContext } from './StyleSheetManager';
+import { DefaultTheme, useTheme } from './ThemeProvider';
 
 const identifiers: { [key: string]: number } = {};
 
@@ -54,13 +54,12 @@ function useInjectedStyle<T extends object>(
   isStatic: boolean,
   resolvedAttrs: T
 ) {
-  const styleSheet = useStyleSheet();
-  const stylis = useStylis();
+  const ssc = useStyleSheetContext();
 
   const className = componentStyle.generateAndInjectStyles(
     isStatic ? EMPTY_OBJECT : resolvedAttrs,
-    styleSheet,
-    stylis
+    ssc.styleSheet,
+    ssc.stylis
   );
 
   if (process.env.NODE_ENV !== 'production') useDebugValue(className);
@@ -110,15 +109,16 @@ function useStyledComponentImpl<Target extends WebTarget, Props extends Executio
     target,
   } = forwardedComponent;
 
-  const defaultShouldForwardProp = useShouldForwardProp();
-  const shouldForwardProp = forwardedComponent.shouldForwardProp || defaultShouldForwardProp;
+  const contextTheme = useTheme();
+  const ssc = useStyleSheetContext();
+  const shouldForwardProp = forwardedComponent.shouldForwardProp || ssc.shouldForwardProp;
 
   if (process.env.NODE_ENV !== 'production') useDebugValue(styledComponentId);
 
   // NOTE: the non-hooks version only subscribes to this when !componentStyle.isStatic,
   // but that'd be against the rules-of-hooks. We could be naughty and do it anyway as it
   // should be an immutable value, but behave for now.
-  const theme = determineTheme(props, useContext(ThemeContext), defaultProps) || EMPTY_OBJECT;
+  const theme = determineTheme(props, contextTheme, defaultProps) || EMPTY_OBJECT;
 
   const context: Dict<any> = resolveContext<Props>(componentAttrs, props, theme);
   const elementToBeCreated: WebTarget = context.as || target;

packages/styled-components/src/models/StyledNativeComponent.ts

@@ -1,4 +1,4 @@
-import React, { createElement, Ref, useContext, useMemo } from 'react';
+import React, { createElement, Ref, useMemo } from 'react';
 import type {
   AttrsArg,
   Dict,
@@ -20,7 +20,7 @@ import hoist from '../utils/hoist';
 import isFunction from '../utils/isFunction';
 import isStyledComponent from '../utils/isStyledComponent';
 import merge from '../utils/mixinDeep';
-import { DefaultTheme, ThemeContext } from './ThemeProvider';
+import { DefaultTheme, useTheme } from './ThemeProvider';
 
 function useResolvedAttrs<Props extends object>(
   theme: DefaultTheme = EMPTY_OBJECT,
@@ -66,10 +66,12 @@ function useStyledComponentImpl<
     target,
   } = forwardedComponent;
 
+  const contextTheme = useTheme();
+
   // NOTE: the non-hooks version only subscribes to this when !componentStyle.isStatic,
   // but that'd be against the rules-of-hooks. We could be naughty and do it anyway as it
   // should be an immutable value, but behave for now.
-  const theme = determineTheme(props, useContext(ThemeContext), defaultProps);
+  const theme = determineTheme(props, contextTheme, defaultProps);
 
   const [context, attrs] = useResolvedAttrs(theme || EMPTY_OBJECT, props, componentAttrs);
 

packages/styled-components/src/models/ThemeProvider.tsx

@@ -66,11 +66,15 @@ function mergeTheme(theme: ThemeArgument, outerTheme?: DefaultTheme): DefaultThe
   return outerTheme ? { ...outerTheme, ...theme } : theme;
 }
 
+export function useTheme(): DefaultTheme | undefined {
+  return useContext(ThemeContext);
+}
+
 /**
  * Provide a theme to an entire react component tree via context
  */
 export default function ThemeProvider(props: Props): JSX.Element | null {
-  const outerTheme = useContext(ThemeContext);
+  const outerTheme = useTheme();
   const themeContext = useMemo(
     () => mergeTheme(props.theme, outerTheme),
     [props.theme, outerTheme]

packages/styled-components/src/models/test/ThemeProvider.test.tsx

@@ -1,8 +1,8 @@
 import React from 'react';
-import TestRenderer from 'react-test-renderer';
+import TestRenderer, { ReactTestInstance } from 'react-test-renderer';
 import withTheme from '../../hoc/withTheme';
 import { resetStyled } from '../../test/utils';
-import ThemeProvider from '../ThemeProvider';
+import ThemeProvider, { useTheme } from '../ThemeProvider';
 
 let styled: ReturnType<typeof resetStyled>;
 
@@ -125,3 +125,38 @@ describe('ThemeProvider', () => {
     expect(wrapper.root.findByType(MyDiv).props.theme).toEqual(expected);
   });
 });
+
+describe('useTheme', () => {
+  beforeEach(() => {
+    styled = resetStyled();
+  });
+
+  it('useTheme should get the same theme that is serving ThemeProvider', () => {
+    const mainTheme = { main: 'black' };
+
+    const MyDivOne = withTheme(styled.div``);
+    const MyDivWithThemeOne = withTheme(MyDivOne);
+    const MyDivWithThemeContext = () => {
+      const theme = useTheme();
+      return <div data-theme={theme} />;
+    };
+
+    const wrapper = TestRenderer.create(
+      <div>
+        <ThemeProvider theme={mainTheme}>
+          <React.Fragment>
+            <MyDivWithThemeOne />
+            <MyDivWithThemeContext />
+          </React.Fragment>
+        </ThemeProvider>
+      </div>
+    );
+
+    expect(wrapper.root.findByType(MyDivOne).props.theme).toEqual(mainTheme);
+    expect(
+      (wrapper.root.findByType(MyDivWithThemeContext).children[0] as ReactTestInstance).props[
+        'data-theme'
+      ]
+    ).toEqual(mainTheme);
+  });
+});

packages/styled-components/src/native/index.ts

@@ -2,10 +2,9 @@ import React from 'react';
 import constructWithOptions, { Styled } from '../constructors/constructWithOptions';
 import css from '../constructors/css';
 import withTheme from '../hoc/withTheme';
-import useTheme from '../hooks/useTheme';
 import _InlineStyle from '../models/InlineStyle';
 import _StyledNativeComponent from '../models/StyledNativeComponent';
-import ThemeProvider, { ThemeConsumer, ThemeContext } from '../models/ThemeProvider';
+import ThemeProvider, { ThemeConsumer, ThemeContext, useTheme } from '../models/ThemeProvider';
 import { NativeTarget } from '../types';
 import isStyledComponent from '../utils/isStyledComponent';
 
@@ -46,12 +45,12 @@ const aliases = [
   'VirtualizedList',
 ] as const;
 
-type KnownComponents = typeof aliases[number];
+type KnownComponents = (typeof aliases)[number];
 
 /** Isolates RN-provided components since they don't expose a helper type for this. */
 type RNComponents = {
-  [K in keyof typeof reactNative]: typeof reactNative[K] extends React.ComponentType<any>
-    ? typeof reactNative[K]
+  [K in keyof typeof reactNative]: (typeof reactNative)[K] extends React.ComponentType<any>
+    ? (typeof reactNative)[K]
     : never;
 };