[pigment-css] Handle more scenarios while transforming sx prop

packages/pigment-react/package.json

@@ -51,12 +51,14 @@
     "stylis": "^4.3.1"
   },
   "devDependencies": {
+    "@babel/plugin-syntax-jsx": "^7.23.3",
     "@types/babel__core": "^7.20.5",
     "@types/babel__helper-module-imports": "^7.18.3",
     "@types/babel__helper-plugin-utils": "^7.10.3",
     "@types/chai": "^4.3.12",
     "@types/cssesc": "^3.0.2",
     "@types/lodash": "^4.14.202",
+    "@types/mocha": "^10.0.6",
     "@types/node": "^18.19.21",
     "@types/react": "^18.2.55",
     "@types/stylis": "^4.2.5",
@@ -133,15 +135,6 @@
     }
   },
   "nx": {
-    "targetDefaults": {
-      "build": {
-        "outputs": [
-          "{projectRoot}/build",
-          "{projectRoot}/processors",
-          "{projectRoot}/utils"
-        ]
-      }
-    },
     "targets": {
       "test": {
         "cache": false,
@@ -154,6 +147,13 @@
         "dependsOn": [
           "build"
         ]
+      },
+      "build": {
+        "outputs": [
+          "{projectRoot}/build",
+          "{projectRoot}/processors",
+          "{projectRoot}/utils"
+        ]
       }
     }
   }

packages/pigment-react/src/utils/pre-linaria-plugin.ts

@@ -1,69 +1,82 @@
 import { addNamed } from '@babel/helper-module-imports';
 import { declare } from '@babel/helper-plugin-utils';
-import { sxObjectExtractor } from './sxObjectExtractor';
+import { NodePath } from '@babel/core';
+import * as Types from '@babel/types';
+import { sxPropConverter } from './sxPropConverter';
 
-export const babelPlugin = declare((api) => {
-  api.assertVersion(7);
-  const { types: t } = api;
-  return {
-    name: '@pigmentcss/zero-babel-plugin',
-    visitor: {
-      JSXAttribute(path) {
-        const namePath = path.get('name');
-        const openingElement = path.findParent((p) => p.isJSXOpeningElement());
-        if (
-          !openingElement ||
-          !openingElement.isJSXOpeningElement() ||
-          !namePath.isJSXIdentifier() ||
-          namePath.node.name !== 'sx'
-        ) {
-          return;
-        }
-        const tagName = openingElement.get('name');
-        if (!tagName.isJSXIdentifier()) {
-          return;
-        }
-        const valuePath = path.get('value');
-        if (!valuePath.isJSXExpressionContainer()) {
-          return;
-        }
-        const expressionPath = valuePath.get('expression');
-        if (!expressionPath.isExpression()) {
-          return;
-        }
-        if (!expressionPath.isObjectExpression() && !expressionPath.isArrowFunctionExpression()) {
-          return;
-        }
-        sxObjectExtractor(expressionPath);
-        const sxIdentifier = addNamed(namePath, 'sx', process.env.PACKAGE_NAME as string);
-        expressionPath.replaceWith(
-          t.callExpression(sxIdentifier, [expressionPath.node, t.identifier(tagName.node.name)]),
-        );
-      },
-      ObjectProperty(path) {
-        // @TODO - Maybe add support for React.createElement calls as well.
-        // Right now, it only checks for jsx(),jsxs(),jsxDEV() and jsxsDEV() calls.
-        const keyPath = path.get('key');
-        if (!keyPath.isIdentifier() || keyPath.node.name !== 'sx') {
-          return;
-        }
-        const valuePath = path.get('value');
-        if (!valuePath.isObjectExpression() && !valuePath.isArrowFunctionExpression()) {
-          return;
-        }
-        const parentJsxCall = path.findParent((p) => p.isCallExpression());
-        if (!parentJsxCall || !parentJsxCall.isCallExpression()) {
-          return;
-        }
-        const callee = parentJsxCall.get('callee');
-        if (!callee.isIdentifier() || !callee.node.name.includes('jsx')) {
-          return;
-        }
-        const jsxElement = parentJsxCall.get('arguments')[0];
-        sxObjectExtractor(valuePath);
-        const sxIdentifier = addNamed(keyPath, 'sx', process.env.PACKAGE_NAME as string);
-        valuePath.replaceWith(t.callExpression(sxIdentifier, [valuePath.node, jsxElement.node]));
-      },
-    },
+function replaceNodePath(
+  expressionPath: NodePath<Types.Expression>,
+  namePath: NodePath<Types.JSXIdentifier | Types.Identifier>,
+  importName: string,
+  t: typeof Types,
+  tagName: NodePath<Types.JSXIdentifier | Types.Identifier>,
+) {
+  const sxIdentifier = addNamed(namePath, importName, process.env.PACKAGE_NAME as string);
+
+  const wrapWithSxCall = (expPath: NodePath<Types.Expression>) => {
+    expPath.replaceWith(
+      t.callExpression(sxIdentifier, [expPath.node, t.identifier(tagName.node.name)]),
+    );
   };
-});
+
+  sxPropConverter(expressionPath, wrapWithSxCall);
+}
+
+export const babelPlugin = declare<{ propName?: string; importName?: string }>(
+  (api, { propName = 'sx', importName = 'sx' }) => {
+    api.assertVersion(7);
+    const { types: t } = api;
+    return {
+      name: '@pigmentcss/zero-babel-plugin',
+      visitor: {
+        JSXAttribute(path) {
+          const namePath = path.get('name');
+          const openingElement = path.findParent((p) => p.isJSXOpeningElement());
+          if (
+            !openingElement ||
+            !openingElement.isJSXOpeningElement() ||
+            !namePath.isJSXIdentifier() ||
+            namePath.node.name !== propName
+          ) {
+            return;
+          }
+          const tagName = openingElement.get('name');
+          if (!tagName.isJSXIdentifier()) {
+            return;
+          }
+          const valuePath = path.get('value');
+          if (!valuePath.isJSXExpressionContainer()) {
+            return;
+          }
+          const expressionPath = valuePath.get('expression');
+          if (!expressionPath.isExpression()) {
+            return;
+          }
+          replaceNodePath(expressionPath, namePath, importName, t, tagName);
+        },
+        ObjectProperty(path) {
+          // @TODO - Maybe add support for React.createElement calls as well.
+          // Right now, it only checks for jsx(),jsxs(),jsxDEV() and jsxsDEV() calls.
+          const keyPath = path.get('key');
+          if (!keyPath.isIdentifier() || keyPath.node.name !== propName) {
+            return;
+          }
+          const valuePath = path.get('value');
+          if (!valuePath.isObjectExpression() && !valuePath.isArrowFunctionExpression()) {
+            return;
+          }
+          const parentJsxCall = path.findParent((p) => p.isCallExpression());
+          if (!parentJsxCall || !parentJsxCall.isCallExpression()) {
+            return;
+          }
+          const callee = parentJsxCall.get('callee');
+          if (!callee.isIdentifier() || !callee.node.name.includes('jsx')) {
+            return;
+          }
+          const jsxElement = parentJsxCall.get('arguments')[0] as NodePath<Types.Identifier>;
+          replaceNodePath(valuePath, keyPath, importName, t, jsxElement);
+        },
+      },
+    };
+  },
+);

packages/pigment-react/src/utils/sxObjectExtractor.ts

@@ -23,7 +23,9 @@ function validateObjectKey(
     return;
   }
   if (!parentCall) {
-    throw keyPath.buildCodeFrameError('Expressions in css object keys are not supported.');
+    throw keyPath.buildCodeFrameError(
+      `${process.env.PACKAGE_NAME}: Expressions in css object keys are not supported.`,
+    );
   }
   if (
     !identifiers.every((item) => {
@@ -41,7 +43,7 @@ function validateObjectKey(
     })
   ) {
     throw keyPath.buildCodeFrameError(
-      'Variables in css object keys should only use the passed theme(s) object or variables that are defined in the root scope.',
+      `${process.env.PACKAGE_NAME}: Variables in css object keys should only use the passed theme(s) object or variables that are defined in the root scope.`,
     );
   }
 }
@@ -59,14 +61,14 @@ function traverseObjectExpression(
       const value = property.get('value');
       if (!value.isExpression()) {
         throw value.buildCodeFrameError(
-          'This value is not supported. It can only be static values or local variables.',
+          `${process.env.PACKAGE_NAME}: This value is not supported. It can only be static values or local variables.`,
         );
       }
       if (value.isObjectExpression()) {
         traverseObjectExpression(value, parentCall);
       } else if (value.isArrowFunctionExpression()) {
         throw value.buildCodeFrameError(
-          'Arrow functions are not supported as values of sx object.',
+          `${process.env.PACKAGE_NAME}: Arrow functions are not supported as values of sx object.`,
         );
       } else if (!value.isLiteral() && !isStaticObjectOrArrayExpression(value)) {
         const identifiers = findIdentifiers([value], 'reference');
@@ -86,7 +88,7 @@ function traverseObjectExpression(
             localIdentifiers.push(id);
           } else {
             throw id.buildCodeFrameError(
-              'Consider moving this variable to the root scope if it has all static values.',
+              `${process.env.PACKAGE_NAME}: Consider moving this variable to the root scope if it has all static values.`,
             );
           }
         });
@@ -103,20 +105,23 @@ function traverseObjectExpression(
       if (
         !identifiers.every((id) => {
           const binding = property.scope.getBinding(id.node.name);
-          if (!binding || binding.scope !== rootScope) {
-            return false;
-          }
-          return true;
+          // the indentifier definition should either be in the root scope or in the same scope
+          // as the object property, ie, ({theme}) => ({...theme.applyStyles()})
+          return binding && (binding.scope === rootScope || binding.scope === property.scope);
         })
       ) {
         throw property.buildCodeFrameError(
-          'You can only use variables that are defined in the root scope of the file.',
+          `${process.env.PACKAGE_NAME}: You can only use variables in the spread that are defined in the root scope of the file.`,
         );
       }
     } else if (property.isObjectMethod()) {
-      throw property.buildCodeFrameError('sx prop object does not support ObjectMethods.');
+      throw property.buildCodeFrameError(
+        `${process.env.PACKAGE_NAME}: sx prop object does not support ObjectMethods.`,
+      );
     } else {
-      throw property.buildCodeFrameError('Unknown property in object.');
+      throw property.buildCodeFrameError(
+        `${process.env.PACKAGE_NAME}: Unknown property in object.`,
+      );
     }
   });
 }
@@ -128,7 +133,7 @@ export function sxObjectExtractor(nodePath: NodePath<ObjectExpression | ArrowFun
     const body = nodePath.get('body');
     if (!body.isObjectExpression()) {
       throw body.buildCodeFrameError(
-        "sx prop only supports arrow functions that directly return an object, e.g. () => ({color: 'red'}). You can accept theme object in the params if required.",
+        `${process.env.PACKAGE_NAME}: sx prop only supports arrow functions directly returning an object, e.g. () => ({color: 'red'}). You can accept theme object in the params if required.`,
       );
     }
     traverseObjectExpression(body, nodePath);

packages/pigment-react/src/utils/sxPropConverter.ts

@@ -0,0 +1,45 @@
+import { NodePath } from '@babel/core';
+import { ArrowFunctionExpression, Expression, ObjectExpression } from '@babel/types';
+import { sxObjectExtractor } from './sxObjectExtractor';
+
+function isAllowedExpression(
+  node: NodePath<Expression>,
+): node is NodePath<ObjectExpression> | NodePath<ArrowFunctionExpression> {
+  return node.isObjectExpression() || node.isArrowFunctionExpression();
+}
+
+export function sxPropConverter(
+  node: NodePath<Expression>,
+  wrapWithSxCall: (expPath: NodePath<Expression>) => void,
+) {
+  if (node.isConditionalExpression()) {
+    const consequent = node.get('consequent');
+    const alternate = node.get('alternate');
+
+    if (isAllowedExpression(consequent)) {
+      sxObjectExtractor(consequent);
+      wrapWithSxCall(consequent);
+    }
+    if (isAllowedExpression(alternate)) {
+      sxObjectExtractor(alternate);
+      wrapWithSxCall(alternate);
+    }
+  } else if (node.isLogicalExpression()) {
+    const right = node.get('right');
+    if (isAllowedExpression(right)) {
+      sxObjectExtractor(right);
+      wrapWithSxCall(right);
+    }
+  } else if (isAllowedExpression(node)) {
+    sxObjectExtractor(node);
+    wrapWithSxCall(node);
+  } else if (node.isIdentifier()) {
+    const rootScope = node.scope.getProgramParent();
+    const binding = node.scope.getBinding(node.node.name);
+    // Simplest case, ie, const styles = {static object}
+    // and is used as <Component sx={styles} />
+    if (binding?.scope === rootScope) {
+      wrapWithSxCall(node);
+    }
+  }
+}

packages/pigment-react/tests/styled/fixtures/styled.input.js

@@ -10,24 +10,24 @@ const rotateKeyframe = keyframes({
 });
 
 const Component = styled.div(({ theme }) => ({
-  color: theme.palette.primary.main,
+  color: (theme.vars ?? theme).palette.primary.main,
   animation: `${rotateKeyframe} 2s ease-out 0s infinite`,
 }));
 
-const SliderRail = styled('span', {
+export const SliderRail = styled('span', {
   name: 'MuiSlider',
   slot: 'Rail',
 })`
-  display: none;
+  display: block;
   position: absolute;
   border-radius: inherit;
   background-color: currentColor;
   opacity: 0.38;
-  font-size: ${({ theme }) => theme.size.font.h1};
+  font-size: ${({ theme }) => (theme.vars ?? theme).size.font.h1};
 `;
 
 const SliderRail2 = styled.span`
   display: block;
   opacity: 0.38;
-  font-size: ${({ theme }) => theme.size.font.h1};
+  font-size: ${({ theme }) => (theme.vars ?? theme).size.font.h1};
 `;

packages/pigment-react/tests/styled/fixtures/styled.output.css

@@ -11,7 +11,7 @@
   animation: r1419f2q 2s ease-out 0s infinite;
 }
 .s1sjy0ja {
-  display: none;
+  display: block;
   position: absolute;
   border-radius: inherit;
   background-color: currentColor;

packages/pigment-react/tests/styled/fixtures/styled.output.js

@@ -5,7 +5,7 @@ import _theme from '@pigment-css/react/theme';
 const Component = /*#__PURE__*/ _styled('div')({
   classes: ['c1vtarpi'],
 });
-const SliderRail = /*#__PURE__*/ _styled2('span', {
+export const SliderRail = /*#__PURE__*/ _styled2('span', {
   name: 'MuiSlider',
   slot: 'Rail',
 })({