Merge pull request #877 from iNerdStack/main

onlook-dev · Dec 11, 2024 · 2d01986 · 2d01986
2 parents 05ff3d2 + 90bb39f
commit 2d01986
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 31 deletions.
diff --git a/apps/studio/electron/main/run/helpers.ts b/apps/studio/electron/main/run/helpers.ts
@@ -5,6 +5,7 @@ import * as fs from 'fs';
 import { customAlphabet } from 'nanoid/non-secure';
 import * as nodePath from 'path';
 import { VALID_DATA_ATTR_CHARS } from '/common/helpers/ids';
+import type { NodePath } from '@babel/traverse';
 
 export const ALLOWED_EXTENSIONS = ['.jsx', '.tsx'];
 export const IGNORED_DIRECTORIES = ['node_modules', 'dist', 'build', '.next', '.git'];
@@ -91,3 +92,30 @@ function getTemplateTag(element: any): TemplateTag {
         },
     };
 }
+
+export function isNodeElementArray(node: t.CallExpression): boolean {
+    return (
+        t.isMemberExpression(node.callee) &&
+        t.isIdentifier(node.callee.property) &&
+        node.callee.property.name === 'map'
+    );
+}
+
+export function getDynamicTypeInfo(path: NodePath<t.JSXElement>): DynamicType | undefined {
+    const parent = path.parent;
+    const grandParent = path.parentPath?.parent;
+
+    // Check for conditional root element
+    const isConditionalRoot =
+        (t.isConditionalExpression(parent) || t.isLogicalExpression(parent)) &&
+        t.isJSXExpressionContainer(grandParent);
+
+    // Check for array map root element
+    const isArrayMapRoot =
+        t.isArrowFunctionExpression(parent) ||
+        (t.isJSXFragment(parent) && path.parentPath?.parentPath?.isArrowFunctionExpression());
+
+    const dynamicType = isConditionalRoot ? 'conditional' : isArrayMapRoot ? 'array' : undefined;
+
+    return dynamicType;
+}
diff --git a/apps/studio/electron/main/run/setup.ts b/apps/studio/electron/main/run/setup.ts
@@ -5,7 +5,14 @@ import type { DynamicType, TemplateNode } from '@onlook/models/element';
 import { generateCode } from '../code/diff/helpers';
 import { formatContent, readFile } from '../code/files';
 import { parseJsxFile } from '../code/helpers';
-import { generateCodeOptions, generateId, getTemplateNode, isReactFragment } from './helpers';
+import {
+    generateCodeOptions,
+    generateId,
+    getTemplateNode,
+    isReactFragment,
+    getDynamicTypeInfo,
+    isNodeElementArray,
+} from './helpers';
 
 export async function getFileWithIds(filePath: string): Promise<string | null> {
     const content = await readFile(filePath);
@@ -68,7 +75,6 @@ function createMapping(ast: t.File, filename: string): Record<string, TemplateNo
     const mapping: Record<string, TemplateNode> = {};
     const componentStack: string[] = [];
     const dynamicTypeStack: DynamicType[] = [];
-    const isFirstLayer: boolean[] = [];
 
     traverse(ast, {
         FunctionDeclaration: {
@@ -97,23 +103,33 @@ function createMapping(ast: t.File, filename: string): Record<string, TemplateNo
         },
         CallExpression: {
             enter(path) {
-                if (
-                    t.isMemberExpression(path.node.callee) &&
-                    t.isIdentifier(path.node.callee.property) &&
-                    path.node.callee.property.name === 'map'
-                ) {
+                if (isNodeElementArray(path.node)) {
                     dynamicTypeStack.push('array');
-                    isFirstLayer.push(true);
                 }
             },
             exit(path) {
-                if (
-                    t.isMemberExpression(path.node.callee) &&
-                    t.isIdentifier(path.node.callee.property) &&
-                    path.node.callee.property.name === 'map'
-                ) {
+                if (isNodeElementArray(path.node)) {
+                    dynamicTypeStack.pop();
+                }
+            },
+        },
+        ConditionalExpression: {
+            enter() {
+                dynamicTypeStack.push('conditional');
+            },
+            exit() {
+                dynamicTypeStack.pop();
+            },
+        },
+        LogicalExpression: {
+            enter(path) {
+                if (path.node.operator === '&&' || path.node.operator === '||') {
+                    dynamicTypeStack.push('conditional');
+                }
+            },
+            exit(path) {
+                if (path.node.operator === '&&' || path.node.operator === '||') {
                     dynamicTypeStack.pop();
-                    isFirstLayer.pop();
                 }
             },
         },
@@ -130,24 +146,9 @@ function createMapping(ast: t.File, filename: string): Record<string, TemplateNo
             if (idAttr) {
                 const elementId = idAttr.value.value;
 
-                const isInFirstLayer = isFirstLayer[isFirstLayer.length - 1];
-                const currentDynamicType =
-                    dynamicTypeStack.length > 0 && isInFirstLayer
-                        ? dynamicTypeStack[dynamicTypeStack.length - 1]
-                        : undefined;
-                const templateNode = getTemplateNode(
-                    path,
-                    filename,
-                    componentStack,
-                    currentDynamicType,
-                );
-
-                mapping[elementId] = templateNode;
-            }
+                const dynamicType = getDynamicTypeInfo(path);
 
-            // After processing the first JSX element in a map, mark that we're no longer in first layer
-            if (isFirstLayer[isFirstLayer.length - 1]) {
-                isFirstLayer[isFirstLayer.length - 1] = false;
+                mapping[elementId] = getTemplateNode(path, filename, componentStack, dynamicType);
             }
         },
     });