Add import

Author: owjs3901

src/code-impl.ts

@@ -5,23 +5,23 @@ import { exportAssets } from './commands/exportAssets'
 import { exportComponents } from './commands/exportComponents'
 import { getComponentName } from './utils'
 
+const DEVUP_COMPONENTS = [
+  'Center',
+  'VStack',
+  'Flex',
+  'Grid',
+  'Box',
+  'Text',
+  'Image',
+]
+
 export function extractImports(
   componentsCodes: ReadonlyArray<readonly [string, string]>,
 ): string[] {
   const allCode = componentsCodes.map(([_, code]) => code).join('\n')
   const imports = new Set<string>()
 
-  const devupComponents = [
-    'Center',
-    'VStack',
-    'Flex',
-    'Grid',
-    'Box',
-    'Text',
-    'Image',
-  ]
-
-  for (const component of devupComponents) {
+  for (const component of DEVUP_COMPONENTS) {
     const regex = new RegExp(`<${component}[\\s/>]`, 'g')
     if (regex.test(allCode)) {
       imports.add(component)
@@ -35,14 +35,54 @@ export function extractImports(
   return Array.from(imports).sort()
 }
 
+export function extractCustomComponentImports(
+  componentsCodes: ReadonlyArray<readonly [string, string]>,
+): string[] {
+  console.log(componentsCodes)
+  const allCode = componentsCodes.map(([_, code]) => code).join('\n')
+  const customImports = new Set<string>()
+
+  // Find all component usages in JSX: <ComponentName or <ComponentName>
+  const componentUsageRegex = /<([A-Z][a-zA-Z0-9]*)/g
+  const matches = allCode.matchAll(componentUsageRegex)
+  for (const match of matches) {
+    const componentName = match[1]
+    // Skip devup-ui components and components defined in this code
+    if (!DEVUP_COMPONENTS.includes(componentName)) {
+      customImports.add(componentName)
+    }
+  }
+
+  return Array.from(customImports).sort()
+}
+
+function generateImportStatements(
+  componentsCodes: ReadonlyArray<readonly [string, string]>,
+): string {
+  const devupImports = extractImports(componentsCodes)
+  const customImports = extractCustomComponentImports(componentsCodes)
+
+  const statements: string[] = []
+
+  if (devupImports.length > 0) {
+    statements.push(
+      `import { ${devupImports.join(', ')} } from '@devup-ui/react'`,
+    )
+  }
+
+  for (const componentName of customImports) {
+    statements.push(
+      `import { ${componentName} } from '@/components/${componentName}'`,
+    )
+  }
+
+  return statements.length > 0 ? `${statements.join('\n')}\n\n` : ''
+}
+
 function generateBashCLI(
   componentsCodes: ReadonlyArray<readonly [string, string]>,
 ): string {
-  const imports = extractImports(componentsCodes)
-  const importStatement =
-    imports.length > 0
-      ? `import { ${imports.join(', ')} } from '@devup-ui/react'\n\n`
-      : ''
+  const importStatement = generateImportStatements(componentsCodes)
 
   const commands = [
     'mkdir -p src/components',
@@ -60,11 +100,7 @@ function generateBashCLI(
 function generatePowerShellCLI(
   componentsCodes: ReadonlyArray<readonly [string, string]>,
 ): string {
-  const imports = extractImports(componentsCodes)
-  const importStatement =
-    imports.length > 0
-      ? `import { ${imports.join(', ')} } from '@devup-ui/react'\n\n`
-      : ''
+  const importStatement = generateImportStatements(componentsCodes)
 
   const commands = [
     'New-Item -ItemType Directory -Force -Path src\\components | Out-Null',