feat: drastically improve the go to source ast

Author: AlemTuzlak

examples/react/basic/src/button-with-props-only.tsx

@@ -1,3 +1,5 @@
+import { Button } from "./button"
+
 export const ButtonWithProps = (props: { children: React.ReactNode }) => {
-  return <button children={props.children} />
+  return <Button children={props.children} />
 }

examples/react/basic/src/feature.tsx

@@ -1,10 +1,14 @@
 import { Button } from './button'
+import { ButtonWithProps } from './button-with-props-only'
 
 export const Feature = () => {
   return (
     <div>
       <h2>Feature component</h2>
       <Button>Nested</Button>
+      <ButtonWithProps>
+        With props
+      </ButtonWithProps>
     </div>
   )
 }

packages/devtools-vite/package.json

@@ -62,6 +62,7 @@
   "devDependencies": {
     "@types/babel__core": "^7.20.5",
     "@types/babel__generator": "^7.27.0",
-    "@types/babel__traverse": "^7.28.0"
+    "@types/babel__traverse": "^7.28.0",
+    "happy-dom": "^18.0.1"
   }
 }

packages/devtools-vite/src/inject-source.ts

@@ -1,43 +1,131 @@
 import { normalizePath } from 'vite'
 import { gen, parse, t, trav } from './babel'
-import type { types as Babel } from '@babel/core'
+import type { types as Babel, NodePath } from '@babel/core'
 import type { ParseResult } from '@babel/parser'
 
+const getPropsNameFromFunctionDeclaration = (functionDeclaration: t.VariableDeclarator | t.FunctionDeclaration) => {
+  let propsName: string | null = null;
+  if (functionDeclaration.type === "FunctionDeclaration") {
+    const firstArgument = functionDeclaration.params[0];
+    // handles (props) => {}
+    if (firstArgument && firstArgument.type === 'Identifier') {
+      propsName = firstArgument.name;
+    }
+    // handles ({ ...props }) => {}
+    if (firstArgument && firstArgument.type === "ObjectPattern") {
+      firstArgument.properties.forEach(prop => {
+        if (prop.type === "RestElement" && prop.argument.type === "Identifier") {
+          propsName = prop.argument.name;
+        }
+      })
+    }
+    return propsName;
+  }
+  // Arrow function case
+  if (functionDeclaration.init?.type === "ArrowFunctionExpression" || functionDeclaration.init?.type === "FunctionExpression") {
+    const firstArgument = functionDeclaration.init.params[0]
+    // handles (props) => {}
+    if (firstArgument && firstArgument.type === 'Identifier') {
+      propsName = firstArgument.name;
+    }
+    // handles ({ ...props }) => {}
+    if (firstArgument && firstArgument.type === "ObjectPattern") {
+      firstArgument.properties.forEach(prop => {
+        if (prop.type === "RestElement" && prop.argument.type === "Identifier") {
+          propsName = prop.argument.name;
+        }
+      })
+    }
+  }
+  return propsName;
+}
+
+const transformJSX = (element: NodePath<t.JSXOpeningElement>, propsName: string | null, file: string) => {
+  const loc = element.node.loc
+  if (!loc) return
+  const line = loc.start.line
+  const column = loc.start.column
+
+  // Check if props are spread and element name starts with lowercase
+  const hasSpread = element.node.attributes.some(
+    (attr) => attr.type === 'JSXSpreadAttribute' && attr.argument.type === "Identifier" && attr.argument.name === propsName,
+  )
+
+  if (hasSpread) {
+    // Do not inject if props are spread and element is lowercase (native HTML)
+    return
+  }
+  if (propsName) {
+    // inject data-source either via props or default to a string "<file>:<line>:<column>"
+    // Inject data-source via props
+    element.node.attributes.push(
+      t.jsxAttribute(
+        t.jsxIdentifier('data-tsd-source'),
+        t.jsxExpressionContainer(
+          t.logicalExpression(
+            '??',
+            t.memberExpression(
+              t.identifier(propsName),
+              t.stringLiteral('data-tsd-source'),
+              true,
+            ),
+            t.stringLiteral(`${file}:${line}:${column}`),
+          ),
+        ),
+      ),
+    )
+
+    return true
+  }
+
+  // Inject data-source as a string: "<file>:<line>:<column>"
+  element.node.attributes.push(
+    t.jsxAttribute(
+      t.jsxIdentifier('data-tsd-source'),
+      t.stringLiteral(`${file}:${line}:${column}`),
+    ),
+  )
+
+
+  return true
+}
+
+
 const transform = (ast: ParseResult<Babel.File>, file: string) => {
   let didTransform = false
   trav(ast, {
-    JSXOpeningElement(path) {
-      const loc = path.node.loc
-      if (!loc) return
-      const line = loc.start.line
-      const column = loc.start.column
-
-      // Check if props are spread and element name starts with lowercase
-      const hasSpread = path.node.attributes.some(
-        (attr) => attr.type === 'JSXSpreadAttribute',
-      )
-      let isLowercase = false
-      const nameNode = path.node.name
-      if (nameNode.type === 'JSXIdentifier') {
-        isLowercase = /^[a-z]/.test(nameNode.name)
-      }
-
-      if (hasSpread && isLowercase) {
-        // Do not inject if props are spread and element is lowercase (native HTML)
+    FunctionDeclaration(functionDeclaration) {
+      const propsName = getPropsNameFromFunctionDeclaration(functionDeclaration.node)
+      functionDeclaration.traverse({
+        JSXOpeningElement(element) {
+          const transformed = transformJSX(element, propsName, file)
+          if (transformed) {
+            didTransform = true
+          }
+        }
+      })
+    },
+    VariableDeclaration(path) {
+      const functionDeclaration = path.node.declarations.find(decl => {
+        return decl.init?.type === "ArrowFunctionExpression" || decl.init?.type === "FunctionExpression"
+      })
+      if (!functionDeclaration) {
         return
       }
+      const propsName = getPropsNameFromFunctionDeclaration(functionDeclaration)
 
-      // Inject data-source as a string: "<file>:<line>:<column>"
-      path.node.attributes.push(
-        t.jsxAttribute(
-          t.jsxIdentifier('data-tsd-source'),
-          t.stringLiteral(`${file}:${line}:${column}`),
-        ),
-      )
+      path.traverse({
+        JSXOpeningElement(element) {
+          const transformed = transformJSX(element, propsName, file)
+          if (transformed) {
+            didTransform = true
+          }
+        }
+      })
+    }
+  },
 
-      didTransform = true
-    },
-  })
+  )
 
   return didTransform
 }

packages/devtools-vite/tests/index.test.ts

@@ -1,7 +1,7 @@
-import { describe, it } from 'vitest'
+import { describe, expect, it } from 'vitest'
 
 describe('test', () => {
   it('works', () => {
-    // test implementation
+    expect(true).toBe(true)
   })
 })