From fe2a84b7dc8c73e90bdd601acf618b3a77040dbf Mon Sep 17 00:00:00 2001
From: Shota Fuji <pockawoooh@gmail.com>
Date: Sun, 29 Nov 2020 03:15:30 +0000
Subject: [PATCH] fix(addon-docs): render Vue classes in dynamic source

https://github.com/storybookjs/storybook/issues/13326

Both static class (`class="foo"`) and dynamic class (`:class="..."`)
are ignored because Vue treats them as special attribute. They don't
exist in normal vnode.data.attrs nor componentOptions.propsData.
---
 .../frameworks/vue/sourceDecorator.test.ts    | 60 +++++++++++++++++++
 .../src/frameworks/vue/sourceDecorator.ts     | 24 ++++++++
 2 files changed, 84 insertions(+)

diff --git a/addons/docs/src/frameworks/vue/sourceDecorator.test.ts b/addons/docs/src/frameworks/vue/sourceDecorator.test.ts
index b02569d9aa62..c4aa381a3a4c 100644
--- a/addons/docs/src/frameworks/vue/sourceDecorator.test.ts
+++ b/addons/docs/src/frameworks/vue/sourceDecorator.test.ts
@@ -30,6 +30,66 @@ describe('vnodeToString', () => {
     ).toMatchInlineSnapshot(`<button >Button</button>`);
   });
 
+  it('static class', () => {
+    expect(
+      vnodeToString(
+        getVNode({
+          template: `<button class="foo bar">Button</button>`,
+        })
+      )
+    ).toMatchInlineSnapshot(`<button class="foo bar">Button</button>`);
+  });
+
+  it('string dynamic class', () => {
+    expect(
+      vnodeToString(
+        getVNode({
+          template: `<button :class="'foo'">Button</button>`,
+        })
+      )
+    ).toMatchInlineSnapshot(`<button class="foo">Button</button>`);
+  });
+
+  it('non-string dynamic class', () => {
+    expect(
+      vnodeToString(
+        getVNode({
+          template: `<button :class="1">Button</button>`,
+        })
+      )
+    ).toMatchInlineSnapshot(`<button >Button</button>`);
+  });
+
+  it('array dynamic class', () => {
+    expect(
+      vnodeToString(
+        getVNode({
+          template: `<button :class="['foo', null, false, 0]">Button</button>`,
+        })
+      )
+    ).toMatchInlineSnapshot(`<button class="foo">Button</button>`);
+  });
+
+  it('merge dynamic and static classes', () => {
+    expect(
+      vnodeToString(
+        getVNode({
+          template: `<button class="foo" :class="{bar: null, baz: 1}">Button</button>`,
+        })
+      )
+    ).toMatchInlineSnapshot(`<button class="foo baz">Button</button>`);
+  });
+
+  it('object dynamic class', () => {
+    expect(
+      vnodeToString(
+        getVNode({
+          template: `<button :class="{foo: true, bar: false}">Button</button>`,
+        })
+      )
+    ).toMatchInlineSnapshot(`<button class="foo">Button</button>`);
+  });
+
   it('attributes', () => {
     const MyComponent: ComponentOptions<any, any, any> = {
       props: ['propA', 'propB', 'propC', 'propD', 'propE', 'propF', 'propG'],
diff --git a/addons/docs/src/frameworks/vue/sourceDecorator.ts b/addons/docs/src/frameworks/vue/sourceDecorator.ts
index 04975f10420c..6178b60dc66f 100644
--- a/addons/docs/src/frameworks/vue/sourceDecorator.ts
+++ b/addons/docs/src/frameworks/vue/sourceDecorator.ts
@@ -76,6 +76,7 @@ export const sourceDecorator = (storyFn: any, context: StoryContext) => {
 export function vnodeToString(vnode: Vue.VNode): string {
   const attrString = [
     ...(vnode.data?.slot ? ([['slot', vnode.data.slot]] as [string, any][]) : []),
+    ['class', normalizeClassAttribute(vnode)],
     ...(vnode.componentOptions?.propsData ? Object.entries(vnode.componentOptions.propsData) : []),
     ...(vnode.data?.attrs ? Object.entries(vnode.data.attrs) : []),
   ]
@@ -122,6 +123,29 @@ export function vnodeToString(vnode: Vue.VNode): string {
     .join('')}</${tag}>`;
 }
 
+function normalizeClassAttribute(vnode: Vue.VNode): string | undefined {
+  if (!vnode.data || (!vnode.data.staticClass && !vnode.data.class)) {
+    return undefined;
+  }
+
+  let dynamicClass: readonly string[] = [];
+
+  if (typeof vnode.data.class === 'string') {
+    dynamicClass = [vnode.data.class];
+  } else if (vnode.data.class instanceof Array) {
+    dynamicClass = vnode.data.class;
+  } else if (typeof vnode.data.class === 'object') {
+    dynamicClass = Object.entries(vnode.data.class)
+      .filter(([, active]) => !!active)
+      .map(([className]) => className);
+  }
+
+  return (
+    [...(vnode.data.staticClass?.split(' ') ?? []), ...dynamicClass].filter(Boolean).join(' ') ||
+    undefined
+  );
+}
+
 function stringifyAttr(attrName: string, value?: any): string | null {
   if (typeof value === 'undefined' || typeof value === 'function') {
     return null;