From ee1e20c3bb9595c599baaa5bc51b8b335b21a2c9 Mon Sep 17 00:00:00 2001
From: Brijesh Bittu <brijeshb42@gmail.com>
Date: Tue, 9 Jul 2024 12:47:02 +0530
Subject: [PATCH] [vite] Allow package authors to specify packages to include
 in vite

optimization during dev mode.
---
 packages/pigment-css-vite-plugin/src/index.ts | 63 ++++++++++++++++---
 1 file changed, 56 insertions(+), 7 deletions(-)

diff --git a/packages/pigment-css-vite-plugin/src/index.ts b/packages/pigment-css-vite-plugin/src/index.ts
index a79dbb82..3a852375 100644
--- a/packages/pigment-css-vite-plugin/src/index.ts
+++ b/packages/pigment-css-vite-plugin/src/index.ts
@@ -1,5 +1,5 @@
 import * as path from 'node:path';
-import type { Plugin } from 'vite';
+import type { Plugin, UserConfig } from 'vite';
 import {
   preprocessor as basePreprocessor,
   generateTokenCss,
@@ -17,6 +17,14 @@ export interface PigmentOptions extends Omit<VitePluginOptions, 'themeArgs'> {
   theme?: Theme;
 }
 
+type PigmentMeta = {
+  'pigment-css'?: {
+    vite?: {
+      includes: string[];
+    };
+  };
+};
+
 const VIRTUAL_CSS_FILE = `\0zero-runtime-styles.css`;
 const VIRTUAL_THEME_FILE = `\0zero-runtime-theme.js`;
 
@@ -37,6 +45,8 @@ function isZeroRuntimeProcessableFile(fileName: string, transformLibraries: stri
   );
 }
 
+const MATERIAL_WRAPPER_LIB = '@mui/material-pigment-css';
+
 export function pigment(options: PigmentOptions) {
   const {
     theme,
@@ -47,13 +57,13 @@ export function pigment(options: PigmentOptions) {
     css,
     ...rest
   } = options ?? {};
-  const finalTransformLibraries = transformLibraries
-    .concat([process.env.RUNTIME_PACKAGE_NAME as string, '@mui/material-pigment-css'])
-    .map((lib) => lib.split('/').join(path.sep));
+  const defaultLibs = [process.env.RUNTIME_PACKAGE_NAME as string, MATERIAL_WRAPPER_LIB];
+  const allLibs = transformLibraries.concat(defaultLibs);
+  const finalTransformLibraries = allLibs.map((lib) => lib.split('/').join(path.sep));
 
   function injectMUITokensPlugin(): Plugin {
     return {
-      name: 'vite-mui-theme-injection-plugin',
+      name: 'pigment-css-theme-injection-plugin',
       enforce: 'pre',
       resolveId(source) {
         if (finalTransformLibraries.some((lib) => source.includes(`${lib}/styles.css`))) {
@@ -78,7 +88,7 @@ export function pigment(options: PigmentOptions) {
 
   function intermediateBabelPlugin(): Plugin {
     return {
-      name: 'vite-mui-zero-intermediate-plugin',
+      name: 'pigment-css-sx-plugin',
       enforce: 'post',
       async transform(code, id) {
         const [filename] = id.split('?');
@@ -104,6 +114,40 @@ export function pigment(options: PigmentOptions) {
     };
   }
 
+  function updateConfigPlugin(): Plugin {
+    return {
+      name: 'pigment-css-config-plugin',
+      config() {
+        const includes: Array<string> = [];
+        allLibs.forEach((lib) => {
+          try {
+            // eslint-disable-next-line import/no-dynamic-require, global-require
+            const pkg = require(`${lib}/package.json`) as PigmentMeta;
+            const pkgIncludes = pkg['pigment-css']?.vite?.includes ?? [];
+            if (pkgIncludes.length > 0) {
+              includes.push(...pkgIncludes);
+            }
+          } catch (ex) {
+            if (lib !== MATERIAL_WRAPPER_LIB) {
+              console.warn(
+                `${process.env.PACKAGE_NAME}: You haveve specified "${lib}" in "transformLibraries" but it is not installed.`,
+              );
+            }
+          }
+        });
+        const optimizeDeps: UserConfig['optimizeDeps'] = {
+          exclude: allLibs,
+        };
+        if (includes.length) {
+          optimizeDeps.include = includes;
+        }
+        return {
+          optimizeDeps,
+        };
+      },
+    };
+  }
+
   const withRtl = (selector: string, cssText: string) => {
     return basePreprocessor(selector, cssText, css);
   };
@@ -128,7 +172,12 @@ export function pigment(options: PigmentOptions) {
     ...rest,
   });
 
-  return [injectMUITokensPlugin(), transformSx ? intermediateBabelPlugin() : null, zeroPlugin];
+  return [
+    updateConfigPlugin(),
+    injectMUITokensPlugin(),
+    transformSx ? intermediateBabelPlugin() : null,
+    zeroPlugin,
+  ];
 }
 
 export { extendTheme };