From 41ccf0567e0f56835717556586b2fdaf7495b3e6 Mon Sep 17 00:00:00 2001
From: hasparus <hasparus@gmail.com>
Date: Fri, 18 Mar 2022 22:11:11 +0100
Subject: [PATCH] chore(components): fix themed()

---
 packages/mdx/src/index.tsx  | 55 ++++++++++++++++++++++++-------------
 packages/mdx/test/index.tsx | 17 ++----------
 2 files changed, 39 insertions(+), 33 deletions(-)

diff --git a/packages/mdx/src/index.tsx b/packages/mdx/src/index.tsx
index e5f3445c5..2e41ccf2c 100644
--- a/packages/mdx/src/index.tsx
+++ b/packages/mdx/src/index.tsx
@@ -86,21 +86,8 @@ const alias = (n: ThemedComponentName): keyof JSX.IntrinsicElements =>
   isAlias(n) ? aliases[n] : n
 
 export const themed =
-  (key: ThemedComponentName | (string & {})) =>
-  ({ theme, ...rest }: ThemedProps) => {
-    const extraStyles: CSSObject = {}
-
-    if (key === 'th' || key === 'td') {
-      const { align } = rest as DetailedHTMLProps<
-        React.ThHTMLAttributes<HTMLTableHeaderCellElement>,
-        HTMLTableHeaderCellElement
-      >
-
-      if (align !== 'char') extraStyles.textAlign = align
-    }
-
-    return css({ ...get(theme, `styles.${key}`), ...extraStyles })(theme)
-  }
+  (key: ThemedComponentName | (string & {})) => (theme: Theme) =>
+    css(get(theme, `styles.${key}`))(theme)
 
 // opt out of typechecking whenever `as` prop is used
 interface AnyComponentProps extends JSX.IntrinsicAttributes {
@@ -109,8 +96,11 @@ interface AnyComponentProps extends JSX.IntrinsicAttributes {
 
 export interface ThemedComponent<Name extends string> {
   (
-    props: Name extends keyof JSX.IntrinsicElements ? ComponentProps<Name> : {}
+    props: (Name extends keyof JSX.IntrinsicElements
+      ? ComponentProps<Name>
+      : {}) & { css?: CSSObject }
   ): JSX.Element
+  displayName: string
 }
 
 export type ThemedComponentsDict = {
@@ -127,14 +117,30 @@ const createThemedComponent = <Name extends string>(
 ): ThemedComponent<Name> => {
   const variantStyles = themed(variant)
 
-  return (props) => {
+  const component: ThemedComponent<Name> = (props) => {
+    const extraStyles: { textAlign?: 'left' | 'right' | 'center' | 'justify' } =
+      {}
+
+    if (name === 'th' || name === 'td') {
+      const { align } = props as DetailedHTMLProps<
+        React.ThHTMLAttributes<HTMLTableHeaderCellElement>,
+        HTMLTableHeaderCellElement
+      >
+
+      if (align !== 'char') extraStyles.textAlign = align
+    }
+
     const css = (props as any)['css']
 
     return jsx(name, {
       ...props,
-      css: css ? [css, variantStyles] : variantStyles,
+      css: [props.css, variantStyles, extraStyles].filter(Boolean),
     })
   }
+
+  component.displayName = `Themed(${name})`
+
+  return component
 }
 
 interface ThemedDivProps
@@ -167,10 +173,21 @@ const createComponents = (comps: MDXProviderComponents) => {
   // todo: test this behaviour
   componentKeys.forEach((key) => {
     const componentAtKey = comps[key]
+
+    console.log('createComponents', key, componentAtKey)
+
     if (componentAtKey) {
-      next[key] = (props) => jsx(componentAtKey, { ...props, css: themed(key) })
+      const component: ThemedComponent<string> = (props) => {
+        console.log('component render props=', props)
+        return jsx(componentAtKey, { ...props, css: themed(key) })
+      }
+
+      component.displayName = "MdxComponents('" + key + "')"
+
+      next[key] = component
     }
   })
+
   return next
 }
 
diff --git a/packages/mdx/test/index.tsx b/packages/mdx/test/index.tsx
index 1501a5032..c88cdd56e 100644
--- a/packages/mdx/test/index.tsx
+++ b/packages/mdx/test/index.tsx
@@ -17,8 +17,10 @@ expect.extend(matchers)
 
 test('styles React components', () => {
   const Beep = (props: React.ComponentPropsWithoutRef<'h2'>) => (
+    // eslint-disable-next-line jsx-a11y/heading-has-content
     <h2 {...props} />
   )
+
   const Inner = (props: React.ComponentPropsWithoutRef<typeof Beep>) =>
     mdx('Beep', props)
 
@@ -45,22 +47,9 @@ test('styles React components', () => {
   expect(json).toHaveStyleRule('color', 'tomato')
 })
 
-test('cleans up style props', () => {
-  const json = renderJSON(
-    // @ts-expect-error
-    <Themed.h1 mx={2} id="test">
-      Hello
-    </Themed.h1>
-  )!
-  expect(json.props.id).toBe('test')
-  expect(json.props.mx).not.toBeDefined()
-})
-
 test('themed extracts styles from the theme', () => {
   expect(
-    themed('footer')({
-      theme: { styles: { footer: { background: 'skyblue' } } },
-    })
+    themed('footer')({ styles: { footer: { background: 'skyblue' } } })
   ).toStrictEqual({ background: 'skyblue' })
 })