Support array for sx prop (#52)

mui · May 6, 2024 · bfca706 · bfca706
1 parent 7a75d3c
commit bfca706
Show file tree

Hide file tree

Showing 22 changed files with 605 additions and 15 deletions.
diff --git a/packages/pigment-css-react/src/processors/sx.ts b/packages/pigment-css-react/src/processors/sx.ts
@@ -145,13 +145,22 @@ export class SxProcessor extends BaseProcessor {
     /**
      * Replace the sx prop with runtime sx
      */
+    let pathToReplace: undefined | NodePath<CallExpression>;
     this.replacer((_tagPath) => {
       const tagPath = _tagPath as NodePath<CallExpression>;
 
-      spreadSxProp(tagPath);
+      const isArrayArgument = spreadSxProp(tagPath);
+      if (isArrayArgument) {
+        pathToReplace = tagPath;
+      }
 
       return tagPath.node;
     }, false);
+
+    if (pathToReplace) {
+      // need to replace outside of `this.replacer` to preserve the import statement
+      pathToReplace.replaceWith(pathToReplace.node.arguments[0]);
+    }
   }
 
   get asSelector(): string {

diff --git a/packages/pigment-css-react/src/sx.d.ts b/packages/pigment-css-react/src/sx.d.ts
@@ -3,4 +3,4 @@ import type { ThemeArgs } from './theme';
 
 export type SxProp = CSSObjectNoCallback | ((themeArgs: ThemeArgs) => CSSObjectNoCallback);
 
-export default function sx(arg: SxProp, componentClass?: string): string;
+export default function sx(arg: SxProp | Array<SxProp>, componentClass?: string): string;
diff --git a/packages/pigment-css-react/src/sx.js b/packages/pigment-css-react/src/sx.js
@@ -1,7 +1,25 @@
 export default function sx(transformedSx, { className, style }) {
-  const sxClass = typeof transformedSx === 'string' ? transformedSx : transformedSx?.className;
-  const sxVars =
-    transformedSx && typeof transformedSx !== 'string' ? transformedSx.vars : undefined;
+  let sxClass = '';
+  let sxVars = {};
+
+  function iterateSx(element) {
+    if (element) {
+      sxClass += `${typeof element === 'string' ? element : element.className} `;
+      sxVars = {
+        ...sxVars,
+        ...(element && typeof element !== 'string' ? element.vars : undefined),
+      };
+    }
+  }
+
+  if (Array.isArray(transformedSx)) {
+    transformedSx.forEach((element) => {
+      iterateSx(element);
+    });
+  } else {
+    iterateSx(transformedSx);
+  }
+
   const varStyles = {};
 
   if (sxVars) {
@@ -15,7 +33,7 @@ export default function sx(transformedSx, { className, style }) {
   }
 
   return {
-    className: `${sxClass}${className ? ` ${className}` : ''}`,
+    className: `${sxClass.trim()}${className ? ` ${className}` : ''}`,
     style: {
       ...varStyles,
       ...style,

diff --git a/packages/pigment-css-react/src/utils/cssFunctionTransformerPlugin.ts b/packages/pigment-css-react/src/utils/cssFunctionTransformerPlugin.ts
@@ -54,7 +54,9 @@ const cssFunctionTransformerPlugin = declare<BabelPluginOptions>((api, pluginOpt
         }
         const themeKeyArr = val.split('.').join('-');
         path.replaceWith(
-          t.stringLiteral(`var(--${finalPrefix}${propertyThemeKey}-${themeKeyArr})`),
+          t.stringLiteral(
+            `var(--${finalPrefix ? `${finalPrefix}-` : ''}${propertyThemeKey}-${themeKeyArr})`,
+          ),
         );
       },
     },

diff --git a/packages/pigment-css-react/src/utils/spreadSxProp.ts b/packages/pigment-css-react/src/utils/spreadSxProp.ts
@@ -1,5 +1,6 @@
 import { NodePath, types as astService } from '@babel/core';
 import {
+  ArrayExpression,
   CallExpression,
   Expression,
   JSXAttribute,
@@ -100,7 +101,7 @@ export default function spreadSxProp(tagPath: NodePath<CallExpression>) {
     | NodePath<ObjectExpression>
     | null;
   if (!target) {
-    return;
+    return false;
   }
   let paths:
     | NodePath<JSXAttribute | JSXSpreadAttribute>[]
@@ -112,19 +113,45 @@ export default function spreadSxProp(tagPath: NodePath<CallExpression>) {
     paths = target.get('properties');
   }
   const { props, sxPath } = getProps(paths);
+  let spreadSxNode: undefined | SpreadElement | JSXSpreadAttribute;
   if (sxPath) {
     const expression = sxPath.get('value');
     if ('node' in expression) {
       if (target.isObjectExpression()) {
-        target.node.properties.push(astService.spreadElement(expression.node as Expression));
+        spreadSxNode = astService.spreadElement(expression.node as Expression);
+        target.node.properties.push(spreadSxNode);
       }
       if (target.isJSXOpeningElement() && expression.isJSXExpressionContainer()) {
-        target.node.attributes.push(
-          astService.jsxSpreadAttribute(expression.node.expression as Expression),
-        );
+        spreadSxNode = astService.jsxSpreadAttribute(expression.node.expression as Expression);
+        target.node.attributes.push(spreadSxNode);
       }
     }
     sxPath.remove();
   }
   tagPath.node.arguments.push(astService.objectExpression(props));
+
+  if (spreadSxNode?.argument.type === 'ArrayExpression') {
+    spreadSxNode.argument = astService.callExpression(tagPath.node.callee, [
+      spreadSxNode.argument,
+      astService.objectExpression(props),
+    ]);
+  }
+
+  // This step is required to pass information about the array argument to the outer function
+  // to replace the `tagPath` with its first argument.
+  //
+  // Check if the sx value is an array expression
+  let arrayPath: NodePath<ArrayExpression> | null = null;
+  if (tagPath.parentPath.isArrayExpression()) {
+    // sx call is a direct child, e.g. [_sx(...), _sx(...)]
+    arrayPath = tagPath.parentPath;
+  } else if (tagPath.parentPath.parentPath?.isArrayExpression()) {
+    // sx call inside a conditional/logical expression, e.g. [true ? _sx(...) : _sx(...), prop && _sx(...)]
+    arrayPath = tagPath.parentPath.parentPath;
+  }
+  if (arrayPath) {
+    return true;
+  }
+
+  return false;
 }
diff --git a/packages/pigment-css-react/src/utils/sx-plugin.ts b/packages/pigment-css-react/src/utils/sx-plugin.ts
@@ -92,6 +92,9 @@ export const babelPlugin = declare<{
         }
         const valuePath = path.get('value');
         if (
+          !valuePath.isIdentifier() &&
+          !valuePath.isMemberExpression() &&
+          !valuePath.isArrayExpression() &&
           !valuePath.isObjectExpression() &&
           !valuePath.isArrowFunctionExpression() &&
           !valuePath.isConditionalExpression() &&

diff --git a/packages/pigment-css-react/src/utils/sxPropConverter.ts b/packages/pigment-css-react/src/utils/sxPropConverter.ts
@@ -5,14 +5,22 @@ import { sxObjectExtractor } from './sxObjectExtractor';
 function isAllowedExpression(
   node: NodePath<Expression>,
 ): node is NodePath<ObjectExpression> | NodePath<ArrowFunctionExpression> {
-  return node.isObjectExpression() || node.isArrowFunctionExpression();
+  return (
+    node.isObjectExpression() || node.isArrowFunctionExpression() || node.isFunctionExpression()
+  );
 }
 
 export function sxPropConverter(
   node: NodePath<Expression>,
   wrapWithSxCall: (expPath: NodePath<Expression>) => void,
 ) {
-  if (node.isConditionalExpression()) {
+  if (node.isArrayExpression()) {
+    node.get('elements').forEach((element) => {
+      if (element.isExpression()) {
+        sxPropConverter(element, wrapWithSxCall);
+      }
+    });
+  } else if (node.isConditionalExpression()) {
     const consequent = node.get('consequent');
     const alternate = node.get('alternate');
 
@@ -41,5 +49,17 @@ export function sxPropConverter(
     if (binding?.scope === rootScope) {
       wrapWithSxCall(node);
     }
+  } else if (node.isMemberExpression()) {
+    let current: NodePath<Expression> = node;
+    while (current.isMemberExpression()) {
+      current = current.get('object');
+    }
+    if (current.isIdentifier()) {
+      const rootScope = current.scope.getProgramParent();
+      const binding = current.scope.getBinding(current.node.name);
+      if (binding?.scope === rootScope) {
+        wrapWithSxCall(node);
+      }
+    }
   }
 }
diff --git a/packages/pigment-css-react/tests/sx/fixtures/sx-array.input.js b/packages/pigment-css-react/tests/sx/fixtures/sx-array.input.js
@@ -0,0 +1,95 @@
+import { jsx as _jsx } from 'react/jsx-runtime';
+
+<div sx={[{ opacity: 1 }]} />;
+
+<div className="foo" style={{ opacity: 1 }} sx={[{ color: 'red' }, { color: 'green' }]} />;
+
+function App(props) {
+  return (
+    <SliderRail
+      {...props}
+      sx={[
+        props.variant === 'secondary'
+          ? { color: props.isRed ? 'red' : 'blue' }
+          : { backgroundColor: 'blue', color: 'white' },
+        (theme) => ({
+          border: `1px solid ${theme.palette.primary.main}`,
+        }),
+      ]}
+    />
+  );
+}
+
+function App2(props) {
+  return (
+    <SliderRail
+      sx={[
+        { color: 'green' },
+        props.variant === 'secondary' && { color: props.isRed ? 'red' : 'blue' },
+      ]}
+      className={`foo ${props.className}`}
+      style={{
+        color: 'red',
+        ...props.style,
+      }}
+    />
+  );
+}
+
+_jsx('div', {
+  sx: [
+    {
+      opacity: 1,
+    },
+  ],
+});
+
+_jsx('div', {
+  className: 'foo',
+  style: {
+    opacity: 1,
+  },
+  sx: [
+    {
+      color: 'red',
+    },
+    {
+      color: 'green',
+    },
+  ],
+});
+
+function App3(props) {
+  return _jsx('div', {
+    sx: [
+      (theme) => ({
+        border: `1px solid ${theme.palette.primary.main}`,
+      }),
+      props.disabled
+        ? {
+            opacity: 0.4,
+          }
+        : {
+            color: 'red',
+          },
+    ],
+    children: 'test',
+    ...props,
+  });
+}
+
+function App4(props) {
+  return _jsx('div', {
+    sx: [
+      (theme) => ({
+        border: `1px solid ${theme.palette.primary.main}`,
+      }),
+      props.variant === 'secondary' && { color: props.isRed ? 'red' : 'blue' },
+    ],
+    className: `foo ${props.className}`,
+    style: {
+      color: 'red',
+      ...props.style,
+    },
+  });
+}
diff --git a/packages/pigment-css-react/tests/sx/fixtures/sx-array.output.css b/packages/pigment-css-react/tests/sx/fixtures/sx-array.output.css
@@ -0,0 +1,49 @@
+.dnf03uj {
+  opacity: 1;
+}
+.dsvwfmp {
+  color: red;
+}
+.d121rcfp {
+  color: green;
+}
+.schfkcb {
+  color: var(--schfkcb-0);
+}
+.s1cp79ho {
+  background-color: blue;
+  color: white;
+}
+.soyt4ij {
+  border: 1px solid red;
+}
+.sr48wd1 {
+  color: green;
+}
+.s6s70bv {
+  color: var(--s6s70bv-0);
+}
+.s1v8upwb {
+  opacity: 1;
+}
+.sjtfjpx {
+  color: red;
+}
+.s1r80n7h {
+  color: green;
+}
+.s1gu7ed8 {
+  border: 1px solid red;
+}
+.s1h4vmh2 {
+  opacity: 0.4;
+}
+.s1oy2sl1 {
+  color: red;
+}
+.s14d8kn5 {
+  border: 1px solid red;
+}
+.s1su4mia {
+  color: var(--s1su4mia-0);
+}