Enable sx prop on any components

apps/pigment-css-next-app/.eslintrc

@@ -3,6 +3,7 @@
   "rules": {
     "import/prefer-default-export": "off",
     "import/extensions": "off",
-    "import/no-unresolved": "off"
-  }
+    "import/no-unresolved": "off",
+    "react/no-unknown-property": ["error", { "ignore": ["sx"] }],
+  },
 }

apps/pigment-css-next-app/next.config.js

@@ -101,7 +101,7 @@ theme.getColorSchemeSelector = (colorScheme) => {
  */
 const pigmentOptions = {
   theme,
-  transformLibraries: ['local-ui-lib'],
+  transformLibraries: ['local-ui-lib', '@mui/material'],
   sourceMap: true,
   displayName: true,
 };

apps/pigment-css-next-app/src/app/page.tsx

@@ -1,7 +1,17 @@
 import Image from 'next/image';
 import { styled, css } from '@pigment-css/react';
+import Button from '@mui/material/Button';
 import styles from './page.module.css';
 
+declare global {
+  namespace React {
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    interface HTMLAttributes<T> {
+      sx?: any;
+    }
+  }
+}
+
 const visuallyHidden = css({
   border: 0,
   clip: 'rect(0 0 0 0)',
@@ -93,8 +103,16 @@ export default function Home() {
   return (
     <Main>
       <div className={visuallyHidden}>I am invisible</div>
-      <Description>
-        <p>
+      <Description
+        sx={{
+          '&::before': {
+            content: '"🚀"',
+            display: 'inline-block',
+            border: '1px solid',
+          },
+        }}
+      >
+        <p sx={{ boxShadow: '0 0 4px 0 rgba(0 0 0 / 0.12)' }}>
           Get started by editing&nbsp;
           <code className={styles.code}>src/app/page.tsx</code>
         </p>
@@ -117,7 +135,12 @@ export default function Home() {
         </div>
       </Description>
 
-      <div className={styles.center}>
+      <div
+        className={styles.center}
+        sx={{
+          border: '1px solid',
+        }}
+      >
         <Image
           className={styles.logo}
           src="/next.svg"
@@ -176,6 +199,18 @@ export default function Home() {
           </h2>
           <p>Instantly deploy your Next.js site to a shareable URL with Vercel.</p>
         </a>
+        <Button
+          sx={{
+            color: 'red',
+            backgroundColor: 'blue',
+            '&:hover': {
+              color: 'blue',
+              backgroundColor: 'red',
+            },
+          }}
+        >
+          Hello
+        </Button>
       </div>
     </Main>
   );

apps/pigment-css-next-app/tsconfig.json

@@ -25,15 +25,6 @@
   "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"],
   "exclude": ["node_modules"],
   "references": [
-    {
-      "path": "../../packages/mui-system/tsconfig.build.json"
-    },
-    {
-      "path": "../../packages/mui-base/tsconfig.build.json"
-    },
-    {
-      "path": "../../packages/mui-material/tsconfig.build.json"
-    },
     {
       "path": "../../packages/pigment-css-react/tsconfig.json"
     }

packages/pigment-css-react/.gitignore

@@ -1,4 +1,3 @@
 /processors/
 /utils/
 LICENSE
-/private-runtime/

packages/pigment-css-react/package.json

@@ -43,9 +43,9 @@
     "@emotion/styled": "^11.11.5",
     "@mui/system": "^6.0.0-alpha.1",
     "@mui/utils": "^6.0.0-alpha.1",
-    "@wyw-in-js/processor-utils": "^0.5.1",
-    "@wyw-in-js/shared": "^0.5.1",
-    "@wyw-in-js/transform": "^0.5.1",
+    "@wyw-in-js/processor-utils": "^0.5.3",
+    "@wyw-in-js/shared": "^0.5.3",
+    "@wyw-in-js/transform": "^0.5.3",
     "clsx": "^2.1.0",
     "cssesc": "^3.0.0",
     "csstype": "^3.1.3",
@@ -91,7 +91,6 @@
   "files": [
     "build",
     "exports",
-    "private-runtime",
     "processors",
     "theme",
     "utils",
@@ -149,15 +148,6 @@
       },
       "require": "./build/RtlProvider.js",
       "default": "./build/RtlProvider.js"
-    },
-    "./private-runtime": {
-      "types": "./private-runtime/index.d.ts",
-      "import": {
-        "types": "./private-runtime/index.d.mts",
-        "default": "./private-runtime/index.mjs"
-      },
-      "require": "./private-runtime/index.js",
-      "default": "./private-runtime/index.js"
     }
   },
   "nx": {

packages/pigment-css-react/src/processors/sx.ts

@@ -1,4 +1,5 @@
-import type { Expression } from '@babel/types';
+import type { NodePath } from '@babel/core';
+import type { CallExpression, Expression } from '@babel/types';
 import {
   validateParams,
   type Params,
@@ -10,31 +11,7 @@ import type { IOptions } from './styled';
 import { processCssObject } from '../utils/processCssObject';
 import { cssFnValueToVariable } from '../utils/cssFnValueToVariable';
 import BaseProcessor from './base-processor';
-
-// @TODO: Maybe figure out a better way allow imports.
-const allowedSxTransformImports = [`${process.env.PACKAGE_NAME}/Box`];
-
-/**
- * Specifically looks for whether the sx prop should be transformed or not.
- * If it's a Pigment CSS styled component, the value of `argumentValue` will
- * be a string className starting with "."
- * In other cases, it explicitly checks if the import is allowed for sx transformation.
- */
-function allowSxTransform(argumentExpression: ExpressionValue, argumentValue?: string) {
-  if (!argumentExpression) {
-    return false;
-  }
-  if (
-    argumentExpression.kind === ValueType.LAZY &&
-    argumentExpression.importedFrom?.some((i) => allowedSxTransformImports.includes(i))
-  ) {
-    return true;
-  }
-  if (argumentValue && argumentValue[0] === '.') {
-    return true;
-  }
-  return false;
-}
+import spreadSxProp from '../utils/spreadSxProp';
 
 export class SxProcessor extends BaseProcessor {
   sxArguments: ExpressionValue[] = [];
@@ -66,10 +43,6 @@ export class SxProcessor extends BaseProcessor {
       }
     }
 
-    if (!allowSxTransform(elementClassExpression, this.elementClassName)) {
-      return;
-    }
-
     let cssText: string = '';
     if (sxStyle.kind === ValueType.CONST) {
       if (sxStyle.ex.type === 'StringLiteral') {
@@ -79,16 +52,13 @@ export class SxProcessor extends BaseProcessor {
       const styleObjOrFn = values.get(sxStyle.ex.name);
       cssText = this.processCss(styleObjOrFn, sxStyle);
     }
-    const selector = this.elementClassName
-      ? `${this.elementClassName}${this.asSelector}`
-      : this.asSelector;
 
     if (!cssText) {
       return;
     }
 
     const rules: Rules = {
-      [selector]: {
+      [this.asSelector]: {
         className: this.className,
         cssText,
         displayName: this.displayName,
@@ -123,6 +93,7 @@ export class SxProcessor extends BaseProcessor {
     if (this.artifacts.length === 0) {
       return;
     }
+    let result = this.value;
     if (this.collectedVariables.length) {
       const varProperties: ReturnType<typeof t.objectProperty>[] = this.collectedVariables.map(
         ([variableId, expression, isUnitLess]) => {
@@ -156,10 +127,31 @@ export class SxProcessor extends BaseProcessor {
         t.objectProperty(t.identifier('className'), t.stringLiteral(this.className)),
         t.objectProperty(t.identifier('vars'), t.objectExpression(varProperties)),
       ]);
-      this.replacer(obj, false);
-    } else {
-      this.replacer(this.value, false);
+      result = obj;
     }
+
+    /**
+     * Replace the sx call with the transformed result. It works for both JSX and non-JSX calls.
+     *
+     * For example:
+     * <Component sx={_sx({ color: 'red' })} /> to <Component sx={_sx('sd5jss7')} />
+     * <Component sx={_sx({ bgcolor: 'red', color: props.color })} /> to <Component sx={_sx({ className: 'bc1d15y', vars: { 'bc1d15y-0': [props.color, false], }})} />
+     */
+    this.replacer((_tagPath) => {
+      const tagPath = _tagPath as NodePath<CallExpression>;
+      return t.callExpression(tagPath.get('callee').node, [result]);
+    }, false);
+
+    /**
+     * Replace the sx prop with runtime sx
+     */
+    this.replacer((_tagPath) => {
+      const tagPath = _tagPath as NodePath<CallExpression>;
+
+      spreadSxProp(tagPath);
+
+      return tagPath.node;
+    }, false);
   }
 
   get asSelector(): string {

packages/pigment-css-react/src/styled.js

@@ -26,7 +26,7 @@ function isHtmlTag(tag) {
   );
 }
 
-const slotShouldForwardProp = (key) => key !== 'sx' && key !== 'as' && key !== 'ownerState';
+const slotShouldForwardProp = (key) => key !== 'as' && key !== 'ownerState';
 const rootShouldForwardProp = (key) => slotShouldForwardProp(key) && key !== 'classes';
 
 /**
@@ -80,7 +80,7 @@ export default function styled(tag, componentMeta = {}) {
     const { displayName, classes = [], vars: cssVars = {}, variants = [] } = options;
 
     const StyledComponent = React.forwardRef(function StyledComponent(inProps, ref) {
-      const { as, className, sx, style, ownerState, ...props } = inProps;
+      const { as, className, style, ownerState, ...props } = inProps;
       const Component = (shouldUseAs && as) || tag;
       const varStyles = Object.entries(cssVars).reduce(
         (acc, [cssVariable, [variableFunction, isUnitLess]]) => {
@@ -97,25 +97,8 @@ export default function styled(tag, componentMeta = {}) {
         },
         {},
       );
-      const sxClass = typeof sx === 'string' ? sx : sx?.className;
-      const sxVars = sx && typeof sx !== 'string' ? sx.vars : undefined;
 
-      if (sxVars) {
-        Object.entries(sxVars).forEach(([cssVariable, [value, isUnitLess]]) => {
-          if (typeof value === 'string' || isUnitLess) {
-            varStyles[`--${cssVariable}`] = value;
-          } else {
-            varStyles[`--${cssVariable}`] = `${value}px`;
-          }
-        });
-      }
-
-      const finalClassName = clsx(
-        classes,
-        sxClass,
-        className,
-        getVariantClasses(inProps, variants),
-      );
+      const finalClassName = clsx(classes, className, getVariantClasses(inProps, variants));
 
       const newProps = {};
       // eslint-disable-next-line no-restricted-syntax

packages/pigment-css-react/src/sx.js

@@ -1,3 +1,24 @@
-export default function sx(styleObj) {
-  return styleObj;
+export default function sx(transformedSx, { className, style }) {
+  const sxClass = typeof transformedSx === 'string' ? transformedSx : transformedSx?.className;
+  const sxVars =
+    transformedSx && typeof transformedSx !== 'string' ? transformedSx.vars : undefined;
+  const varStyles = {};
+
+  if (sxVars) {
+    Object.entries(sxVars).forEach(([cssVariable, [value, isUnitLess]]) => {
+      if (typeof value === 'string' || isUnitLess) {
+        varStyles[`--${cssVariable}`] = value;
+      } else {
+        varStyles[`--${cssVariable}`] = `${value}px`;
+      }
+    });
+  }
+
+  return {
+    className: `${sxClass}${className ? ` ${className}` : ''}`,
+    style: {
+      ...varStyles,
+      ...style,
+    },
+  };
 }