fix reassign forwardProp

mui · Apr 30, 2024 · a5f2b38 · a5f2b38
1 parent 8735a94
commit a5f2b38
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/packages/pigment-css-react/src/styled.js b/packages/pigment-css-react/src/styled.js
@@ -123,6 +123,17 @@ export default function styled(tag, componentMeta = {}) {
         getVariantClasses(inProps, variants),
       );
 
+      if (shouldUseAs && !shouldForwardProp) {
+        // Reassign `shouldForwardProp` based on the incoming `as` prop
+        if (isHtmlTag(Component)) {
+          finalShouldForwardProp = isPropValid;
+        } else if (slot === 'Root' || slot === 'root') {
+          finalShouldForwardProp = rootShouldForwardProp;
+        } else {
+          finalShouldForwardProp = slotShouldForwardProp;
+        }
+      }
+
       const newProps = {};
       // eslint-disable-next-line no-restricted-syntax
       for (const key in props) {

diff --git a/packages/pigment-css-react/tests/styled/runtime-styled.test.js b/packages/pigment-css-react/tests/styled/runtime-styled.test.js
@@ -252,6 +252,20 @@ describe('props filtering', () => {
       expect(container.firstChild).to.have.style('--foo', '300px');
     });
 
+    it('use component forward prop if provided `as` is a component', () => {
+      const StyledDiv = styled('div')({
+        classes: ['root'],
+      });
+
+      function Component({ TagComponent = 'span', ...props }) {
+        return <TagComponent {...props} />;
+      }
+
+      const { container } = render(<StyledDiv as={Component} TagComponent="button" disabled />);
+
+      expect(container.firstChild).to.have.tagName('button');
+    });
+
     it('should forward `as` prop', () => {
       // The components below is a simplified version of the `NativeSelect` component from Material UI.