fix: Nested For transform

vkbansal · Dec 4, 2019 · e81cacc · e81cacc
1 parent 5dd5fda
commit e81cacc
Show file tree

Hide file tree

Showing 6 changed files with 206 additions and 23 deletions.
diff --git a/package.json b/package.json
@@ -6,7 +6,6 @@
   "main": "./lib/index.js",
   "module": "./esm/index.js",
   "esnext": "./esnext/index.js",
-  "typings": "./esm/index.d.ts",
   "scripts": {
     "release": "standard-version",
     "test": "jest src",

diff --git a/src/__tests__/__snapshots__/ForTransformer.test.ts.snap b/src/__tests__/__snapshots__/ForTransformer.test.ts.snap
@@ -7,15 +7,56 @@ Content:
 
 import React from 'react';
 
-import { For } from '@vkbansal/tsx-control-statements';
+import { For, If } from '@vkbansal/tsx-control-statements';
 
 export function ForStatement(): React.ReactElement {
   return (
     <div>
-      <For items={[1, 2, 3, 4, 5]}>{(item) => <div key={item}>{item}</div>}</For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => (
+          <If condition={item > 3}>
+            <div key={item}>{item}</div>
+          </If>
+        )}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => {
+          return (
+            <If condition={item > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => {
+          const a = item;
+
+          return (
+            <If condition={a > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
       <For items={[1, 2, 3, 4, 5]}>
         {function(item) {
-          return <div key={item}>{item}</div>;
+          return (
+            <If condition={item > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {function(item) {
+          const a = item;
+
+          return (
+            <If condition={a > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
         }}
       </For>
     </div>
@@ -26,14 +67,53 @@ export function ForStatement(): React.ReactElement {
 Code before Transform:
 
 import React from 'react';
-import { For } from '@vkbansal/tsx-control-statements';
+import { For, If } from '@vkbansal/tsx-control-statements';
 export function ForStatement(): React.ReactElement {
   return (
     <div>
-      <For items={[1, 2, 3, 4, 5]}>{(item) => <div key={item}>{item}</div>}</For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => (
+          <If condition={item > 3}>
+            <div key={item}>{item}</div>
+          </If>
+        )}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => {
+          return (
+            <If condition={item > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => {
+          const a = item;
+          return (
+            <If condition={a > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
       <For items={[1, 2, 3, 4, 5]}>
         {function(item) {
-          return <div key={item}>{item}</div>;
+          return (
+            <If condition={item > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {function(item) {
+          const a = item;
+          return (
+            <If condition={a > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
         }}
       </For>
     </div>
@@ -47,11 +127,42 @@ import React from 'react';
 export function ForStatement(): React.ReactElement {
   return (
     <div>
-      {[1, 2, 3, 4, 5].map((item) => (
-        <div key={item}>{item}</div>
-      ))}
+      {[1, 2, 3, 4, 5].map((item) =>
+        item > 3 ? (
+          <React.Fragment>
+            <div key={item}>{item}</div>
+          </React.Fragment>
+        ) : null
+      )}
+      {[1, 2, 3, 4, 5].map((item) => {
+        return item > 3 ? (
+          <React.Fragment>
+            <div key={item}>{item}</div>
+          </React.Fragment>
+        ) : null;
+      })}
+      {[1, 2, 3, 4, 5].map((item) => {
+        const a = item;
+        return a > 3 ? (
+          <React.Fragment>
+            <div key={item}>{item}</div>
+          </React.Fragment>
+        ) : null;
+      })}
+      {[1, 2, 3, 4, 5].map(function(item) {
+        return item > 3 ? (
+          <React.Fragment>
+            <div key={item}>{item}</div>
+          </React.Fragment>
+        ) : null;
+      })}
       {[1, 2, 3, 4, 5].map(function(item) {
-        return <div key={item}>{item}</div>;
+        const a = item;
+        return a > 3 ? (
+          <React.Fragment>
+            <div key={item}>{item}</div>
+          </React.Fragment>
+        ) : null;
       })}
     </div>
   );

diff --git a/src/__tests__/fixtures/ForStatement.tsx b/src/__tests__/fixtures/ForStatement.tsx
@@ -1,14 +1,55 @@
 import React from 'react';
 
-import { For } from '@vkbansal/tsx-control-statements';
+import { For, If } from '@vkbansal/tsx-control-statements';
 
 export function ForStatement(): React.ReactElement {
   return (
     <div>
-      <For items={[1, 2, 3, 4, 5]}>{(item) => <div key={item}>{item}</div>}</For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => (
+          <If condition={item > 3}>
+            <div key={item}>{item}</div>
+          </If>
+        )}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => {
+          return (
+            <If condition={item > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {(item) => {
+          const a = item;
+
+          return (
+            <If condition={a > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
       <For items={[1, 2, 3, 4, 5]}>
         {function(item) {
-          return <div key={item}>{item}</div>;
+          return (
+            <If condition={item > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
+        }}
+      </For>
+      <For items={[1, 2, 3, 4, 5]}>
+        {function(item) {
+          const a = item;
+
+          return (
+            <If condition={a > 3}>
+              <div key={item}>{item}</div>
+            </If>
+          );
         }}
       </For>
     </div>

diff --git a/src/transformer/chooseTransformer.ts b/src/transformer/chooseTransformer.ts
@@ -6,7 +6,7 @@ export default function chooseTransformer(
   this: ts.SourceFile,
   node: ts.JsxElement,
   visitor: ts.Visitor
-) {
+): ts.Node {
   const whenElements: ts.JsxElement[] = [];
   let otherwiseElement: ts.JsxElement | null = null;
 

diff --git a/src/transformer/forTransformer.ts b/src/transformer/forTransformer.ts
@@ -1,12 +1,10 @@
 import ts from 'typescript';
 
-// import { getChildrenWrappedInFragmentElement } from './astUtils';
-
 export default function forTransformer(
   this: ts.SourceFile,
   node: ts.JsxElement,
-  _visitor: ts.Visitor
-) {
+  visitor: ts.Visitor
+): ts.Node {
   const items = node.openingElement.attributes.properties.find((attribute) => {
     return attribute.name?.getText(this) === 'items';
   });
@@ -18,20 +16,54 @@ export default function forTransformer(
       (ts.isArrowFunction(child.expression) || ts.isFunctionExpression(child.expression))
   ) as ts.JsxExpression[];
 
+  const firstChild = children[0];
+
   if (
     items &&
     ts.isJsxAttribute(items) &&
     items.initializer &&
     ts.isJsxExpression(items.initializer) &&
-    items.initializer.expression
+    items.initializer.expression &&
+    firstChild &&
+    firstChild.expression
   ) {
     if (children.length > 0) {
-      const firstChild = children[0];
+      let funcExpression = firstChild.expression;
+
+      if (ts.isArrowFunction(funcExpression)) {
+        funcExpression = ts.createArrowFunction(
+          funcExpression.modifiers,
+          funcExpression.typeParameters,
+          funcExpression.parameters,
+          funcExpression.type,
+          funcExpression.equalsGreaterThanToken,
+          ts.isBlock(funcExpression.body)
+            ? ts.createBlock(
+                // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+                funcExpression.body.statements.map(visitor.bind(this)!) as ts.Statement[]
+              )
+            : // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+              (funcExpression.body.forEachChild(visitor.bind(this)!) as ts.ConciseBody)
+        );
+      }
+
+      if (ts.isFunctionExpression(funcExpression)) {
+        funcExpression = ts.createFunctionExpression(
+          funcExpression.modifiers,
+          funcExpression.asteriskToken,
+          funcExpression.name,
+          funcExpression.typeParameters,
+          funcExpression.parameters,
+          funcExpression.type,
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          ts.createBlock(funcExpression.body.statements.map(visitor.bind(this)!) as ts.Statement[])
+        );
+      }
 
       const func = ts.createCall(
         ts.createPropertyAccess(items.initializer.expression, 'map'),
         undefined,
-        [firstChild.expression!]
+        [funcExpression]
       );
       return node.parent && ts.isJsxElement(node.parent)
         ? ts.createJsxExpression(undefined, func)

diff --git a/src/transformer/ifTransformer.ts b/src/transformer/ifTransformer.ts
@@ -6,7 +6,7 @@ export default function ifTransformer(
   this: ts.SourceFile,
   node: ts.JsxElement,
   visitor: ts.Visitor
-) {
+): ts.Node {
   const terinary = createNestedTerinaryExpression.call(this, [node], null, visitor);
 
   return node.parent && ts.isJsxElement(node.parent)