From 775103af37df69d34c79f12c4c1776c47d07f0a0 Mon Sep 17 00:00:00 2001
From: Evan You <evan@vuejs.org>
Date: Thu, 8 Aug 2024 12:35:00 +0800
Subject: [PATCH] feat(custom-element): useHost() helper

---
 .../__tests__/customElement.spec.ts           | 23 ++++++++++++++---
 packages/runtime-dom/src/apiCustomElement.ts  | 25 ++++++++++++-------
 packages/runtime-dom/src/index.ts             |  1 +
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/packages/runtime-dom/__tests__/customElement.spec.ts b/packages/runtime-dom/__tests__/customElement.spec.ts
index c77828b66..eb7a4b368 100644
--- a/packages/runtime-dom/__tests__/customElement.spec.ts
+++ b/packages/runtime-dom/__tests__/customElement.spec.ts
@@ -14,6 +14,7 @@ import {
   ref,
   render,
   renderSlot,
+  useHost,
   useShadowRoot,
 } from '../src'
 
@@ -975,8 +976,22 @@ describe('defineCustomElement', () => {
     })
   })
 
-  describe('useShadowRoot', () => {
-    test('should work for style injection', () => {
+  describe('helpers', () => {
+    test('useHost', () => {
+      const Foo = defineCustomElement({
+        setup() {
+          const host = useHost()!
+          host.setAttribute('id', 'host')
+          return () => h('div', 'hello')
+        },
+      })
+      customElements.define('my-el-use-host', Foo)
+      container.innerHTML = `<my-el-use-host>`
+      const el = container.childNodes[0] as VueElement
+      expect(el.id).toBe('host')
+    })
+
+    test('useShadowRoot for style injection', () => {
       const Foo = defineCustomElement({
         setup() {
           const root = useShadowRoot()!
@@ -986,8 +1001,8 @@ describe('defineCustomElement', () => {
           return () => h('div', 'hello')
         },
       })
-      customElements.define('my-el', Foo)
-      container.innerHTML = `<my-el></my-el>`
+      customElements.define('my-el-use-shadow-root', Foo)
+      container.innerHTML = `<my-el-use-shadow-root>`
       const el = container.childNodes[0] as VueElement
       const style = el.shadowRoot?.querySelector('style')!
       expect(style.textContent).toBe(`div { color: red; }`)
diff --git a/packages/runtime-dom/src/apiCustomElement.ts b/packages/runtime-dom/src/apiCustomElement.ts
index 694f9e0c7..af6063cdb 100644
--- a/packages/runtime-dom/src/apiCustomElement.ts
+++ b/packages/runtime-dom/src/apiCustomElement.ts
@@ -653,24 +653,31 @@ export class VueElement
   }
 }
 
-/**
- * Retrieve the shadowRoot of the current custom element. Only usable in setup()
- * of a `defineCustomElement` component.
- */
-export function useShadowRoot(): ShadowRoot | null {
+export function useHost(caller?: string): VueElement | null {
   const instance = getCurrentInstance()
-  const el = instance && instance.ce
+  const el = instance && (instance.ce as VueElement)
   if (el) {
-    return (el as VueElement).shadowRoot
+    return el
   } else if (__DEV__) {
     if (!instance) {
-      warn(`useShadowRoot called without an active component instance.`)
+      warn(
+        `${caller || 'useHost'} called without an active component instance.`,
+      )
     } else {
       warn(
-        `useShadowRoot can only be used in components defined via ` +
+        `${caller || 'useHost'} can only be used in components defined via ` +
           `defineCustomElement.`,
       )
     }
   }
   return null
 }
+
+/**
+ * Retrieve the shadowRoot of the current custom element. Only usable in setup()
+ * of a `defineCustomElement` component.
+ */
+export function useShadowRoot(): ShadowRoot | null {
+  const el = __DEV__ ? useHost('useShadowRoot') : useHost()
+  return el && el.shadowRoot
+}
diff --git a/packages/runtime-dom/src/index.ts b/packages/runtime-dom/src/index.ts
index 706401ddd..95f27353a 100644
--- a/packages/runtime-dom/src/index.ts
+++ b/packages/runtime-dom/src/index.ts
@@ -247,6 +247,7 @@ export {
   defineCustomElement,
   defineSSRCustomElement,
   useShadowRoot,
+  useHost,
   VueElement,
   type VueElementConstructor,
   type CustomElementOptions,