From a32fac4b26595d7e0f5bbf82e70963367a119c51 Mon Sep 17 00:00:00 2001
From: u3u <qwq@qwq.cat>
Date: Sat, 17 Aug 2019 16:38:56 +0800
Subject: [PATCH] feat(babel-sugar-inject-h): `setup` method support

---
 packages/babel-sugar-inject-h/src/index.js | 10 +++++++++-
 packages/babel-sugar-inject-h/test/test.js | 16 ++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/packages/babel-sugar-inject-h/src/index.js b/packages/babel-sugar-inject-h/src/index.js
index 6784e39..7c35d2c 100644
--- a/packages/babel-sugar-inject-h/src/index.js
+++ b/packages/babel-sugar-inject-h/src/index.js
@@ -1,5 +1,7 @@
 import syntaxJsx from '@babel/plugin-syntax-jsx'
 
+const entry = 'vue-function-api'
+
 /**
  * Check if first parameter is `h`
  * @param t
@@ -62,6 +64,7 @@ export default babel => {
             }
 
             const isRender = path.node.key.name === 'render'
+            const isSetup = path.type === 'ObjectMethod' && path.node.key.name === 'setup'
 
             path
               .get('body')
@@ -70,7 +73,12 @@ export default babel => {
                 t.variableDeclaration('const', [
                   t.variableDeclarator(
                     t.identifier('h'),
-                    isRender
+                    isSetup
+                      ? t.memberExpression(
+                          t.callExpression(t.identifier('require'), [t.stringLiteral(entry)]),
+                          t.identifier('createElement'),
+                        )
+                      : isRender
                       ? t.memberExpression(t.identifier('arguments'), t.numericLiteral(0), true)
                       : t.memberExpression(t.thisExpression(), t.identifier('$createElement')),
                   ),
diff --git a/packages/babel-sugar-inject-h/test/test.js b/packages/babel-sugar-inject-h/test/test.js
index 173c6e2..861b763 100644
--- a/packages/babel-sugar-inject-h/test/test.js
+++ b/packages/babel-sugar-inject-h/test/test.js
@@ -101,6 +101,22 @@ const tests = [
     return <div>test</div>;
   }
 
+};`,
+  },
+  {
+    name: 'Simple injection in object setup methods',
+    from: `const obj = {
+      setup() {
+        return () => <div>test</div>
+      }
+    }`,
+    to: `const obj = {
+  setup() {
+    const h = require("vue-function-api").createElement;
+
+    return () => <div>test</div>;
+  }
+
 };`,
   },
 ]