fix: cli should rewrite the aliases

Author: nmn

.flowconfig

@@ -30,7 +30,7 @@ module.name_mapper='^@stylexjs\/shared\/lib\/\([a-zA-Z0-9_\-]+\)$' -> '<PROJECT_
 module.name_mapper='^@stylexjs/stylex$' -> '<PROJECT_ROOT>/packages/stylex/src/stylex.js'
 module.name_mapper='^@stylexjs/stylex\/lib\/\([a-zA-Z0-9_\-]+\)$' -> '<PROJECT_ROOT>/packages/stylex/src/\1'
 module.name_mapper='^@stylexjs/babel-plugin$' -> '<PROJECT_ROOT>/packages/babel-plugin/src/index.js'
-module.name_mapper='^@stylexjs/babel-plugin\/lib\/\([a-zA-Z0-9_\-]+\)$' -> '<PROJECT_ROOT>/packages/babel-plugin/src/\1'
+module.name_mapper='^@stylexjs/babel-plugin\/lib\/\(.+\)$' -> '<PROJECT_ROOT>/packages/babel-plugin/src/\1'
 ; type-stubs
 module.system.node.resolve_dirname=flow_modules
 module.system.node.resolve_dirname=node_modules

package-lock.json

@@ -671,7 +671,7 @@ const getPossibleFilePaths = (filePath: string) => {
 
 // a function that resolves the absolute path of a file when given the
 // relative path of the file from the source file
-const filePathResolver = (
+export const filePathResolver = (
   relativeFilePath: string,
   sourceFilePath: string,
   aliases: StyleXStateOptions['aliases'],
@@ -719,10 +719,11 @@ const addFileExtension = (
   return importedFilePath + fileExtension;
 };
 
-const matchesFileSuffix = (allowedSuffix: string) => (filename: string) =>
-  ['', ...EXTENSIONS].some((extension) =>
-    filename.endsWith(`${allowedSuffix}${extension}`),
-  );
+export const matchesFileSuffix =
+  (allowedSuffix: string) => (filename: string) =>
+    ['', ...EXTENSIONS].some((extension) =>
+      filename.endsWith(`${allowedSuffix}${extension}`),
+    );
 
 const getProgramPath = (path: NodePath<>): null | NodePath<t.Program> => {
   let programPath = path;

packages/babel-plugin/src/utils/state-manager.js

@@ -19,6 +19,7 @@
     "@babel/plugin-syntax-jsx": "^7.25.7",
     "@babel/plugin-syntax-typescript": "^7.25.7",
     "@babel/types": "^7.25.8",
+    "@dual-bundle/import-meta-resolve": "^4.1.0",
     "@stylexjs/babel-plugin": "0.10.0",
     "ansis": "^3.3.2",
     "fb-watchman": "^2.0.2",

packages/cli/package.json

@@ -8,7 +8,7 @@
  * @flow strict
  */
 
-import type { Rule } from '@stylexjs/babel-plugin';
+import type { Rule, Options as StyleXOptions } from '@stylexjs/babel-plugin';
 
 export type ModuleType =
   | string
@@ -24,7 +24,7 @@ export type CliConfig = {
   babelPluginsPost?: $ReadOnlyArray<any>,
   modules_EXPERIMENTAL: $ReadOnlyArray<ModuleType>,
   useCSSLayers?: boolean,
-  styleXConfig?: { +[string]: mixed },
+  styleXConfig?: StyleXOptions,
 };
 
 export type TransformConfig = {

packages/cli/src/config.js

@@ -8,7 +8,7 @@
  * @flow strict
  */
 
-import type { Rule } from '@stylexjs/babel-plugin';
+import type { Rule, Options as StyleXOptions } from '@stylexjs/babel-plugin';
 import yargs from 'yargs';
 import path from 'path';
 import ansis from 'ansis';
@@ -71,8 +71,7 @@ const babelPresets: $ReadOnlyArray<any> = args.babelPresets;
 const babelPluginsPre: $ReadOnlyArray<any> = args.babelPluginsPre;
 const babelPluginsPost: $ReadOnlyArray<any> = args.babelPluginsPost;
 const useCSSLayers: boolean = args.useCSSLayers;
-const styleXConfig: { +[string]: mixed } =
-  (config.styleXConfig as $FlowFixMe) ?? {};
+const styleXConfig: StyleXOptions = (config.styleXConfig as $FlowFixMe) ?? {};
 
 const cliArgsConfig: CliConfig = {
   input,

packages/cli/src/index.js

@@ -12,12 +12,17 @@ import type { NodePath } from '@babel/traverse';
 import { getRelativePath } from './files';
 import { findModuleDir } from './modules';
 import * as t from '@babel/types';
+import { moduleResolve } from '@dual-bundle/import-meta-resolve';
+import url from 'url';
 
 import * as nodePath from 'path';
 
 type ImportModifierPlugin = $ReadOnly<{
   visitor: {
-    Program: { enter(path: NodePath<t.Program>): void },
+    Program: {
+      enter?: (path: NodePath<t.Program>) => void,
+      exit?: (path: NodePath<t.Program>) => void,
+    },
   },
 }>;
 
@@ -116,3 +121,143 @@ export const createModuleImportModifierPlugin = (
     },
   };
 };
+
+export const createAliasRewritePlugin = (
+  sourceFilePath: string,
+  aliasConfig: $ReadOnly<{ [string]: string | $ReadOnlyArray<string> }>,
+): ImportModifierPlugin => {
+  return {
+    visitor: {
+      Program: {
+        exit(path: NodePath<t.Program>) {
+          path.traverse({
+            ImportDeclaration: {
+              enter(path: NodePath<t.ImportDeclaration>) {
+                const source = path.node.source.value;
+
+                const aliases =
+                  aliasConfig == null
+                    ? aliasConfig
+                    : Object.fromEntries(
+                        Object.entries(aliasConfig).map(([key, value]) => {
+                          if (typeof value === 'string') {
+                            return [key, [value]];
+                          }
+                          return [key, value];
+                        }),
+                      );
+
+                const themeFileExtension = '.stylex';
+                if (!matchesFileSuffix(themeFileExtension)(source)) {
+                  return;
+                }
+                const resolvedFilePath = filePathResolver(
+                  source,
+                  sourceFilePath,
+                  aliases,
+                );
+
+                if (resolvedFilePath == null) {
+                  return;
+                }
+
+                let relativeFilePath = getRelativePath(
+                  sourceFilePath,
+                  resolvedFilePath,
+                );
+
+                const extension = EXTENSIONS.find((ext) =>
+                  relativeFilePath.endsWith(ext),
+                );
+                if (extension != null) {
+                  relativeFilePath = relativeFilePath.slice(
+                    0,
+                    -extension.length,
+                  );
+                }
+
+                path.node.source.value = relativeFilePath;
+              },
+            },
+          });
+        },
+      },
+    },
+  };
+};
+
+const matchesFileSuffix = (allowedSuffix: string) => (filename: string) =>
+  filename.endsWith(`${allowedSuffix}.js`) ||
+  filename.endsWith(`${allowedSuffix}.ts`) ||
+  filename.endsWith(`${allowedSuffix}.tsx`) ||
+  filename.endsWith(`${allowedSuffix}.jsx`) ||
+  filename.endsWith(`${allowedSuffix}.mjs`) ||
+  filename.endsWith(`${allowedSuffix}.cjs`) ||
+  filename.endsWith(allowedSuffix);
+
+const EXTENSIONS = ['.js', '.ts', '.tsx', '.jsx', '.mjs', '.cjs'];
+const filePathResolver = (
+  relativeFilePath: string,
+  sourceFilePath: string,
+  aliases: $ReadOnly<{ [string]: $ReadOnlyArray<string> }>,
+): ?string => {
+  // Try importing without adding any extension
+  // and then every supported extension
+  for (const ext of ['', ...EXTENSIONS]) {
+    const importPathStr = relativeFilePath + ext;
+
+    // Try to resolve relative paths as is
+    if (importPathStr.startsWith('.')) {
+      try {
+        return moduleResolve(importPathStr, url.pathToFileURL(sourceFilePath))
+          .pathname;
+      } catch {
+        continue;
+      }
+    }
+
+    // Otherwise, try to resolve the path with aliases
+    const allAliases = possibleAliasedPaths(importPathStr, aliases);
+    for (const possiblePath of allAliases) {
+      try {
+        return moduleResolve(possiblePath, url.pathToFileURL(sourceFilePath))
+          .pathname;
+      } catch {
+        continue;
+      }
+    }
+  }
+  // Failed to resolve the file path
+  return null;
+};
+
+function possibleAliasedPaths(
+  importPath: string,
+  aliases: $ReadOnly<{ [string]: $ReadOnlyArray<string> }>,
+): $ReadOnlyArray<string> {
+  const result = [importPath];
+  if (aliases == null || Object.keys(aliases).length === 0) {
+    return result;
+  }
+
+  for (const [alias, value] of Object.entries(aliases)) {
+    if (alias.includes('*')) {
+      const [before, after] = alias.split('*');
+      if (importPath.startsWith(before) && importPath.endsWith(after)) {
+        const replacementString = importPath.slice(
+          before.length,
+          after.length > 0 ? -after.length : undefined,
+        );
+        value.forEach((v) => {
+          result.push(v.split('*').join(replacementString));
+        });
+      }
+    } else if (alias === importPath) {
+      value.forEach((v) => {
+        result.push(v);
+      });
+    }
+  }
+
+  return result;
+}

packages/cli/src/plugins.js

@@ -31,6 +31,7 @@ import {
   getDefaultCachePath,
 } from './cache';
 import {
+  createAliasRewritePlugin,
   createImportPlugin,
   createModuleImportModifierPlugin,
 } from './plugins';
@@ -155,6 +156,8 @@ export async function transformFile(
       : path.join(config.output, config.styleXBundleName),
   );
 
+  const aliases = config.styleXConfig?.aliases;
+
   const result = await babel.transformFileAsync(inputFilePath, {
     babelrc: false,
     presets: config.babelPresets,
@@ -173,7 +176,12 @@ export async function transformFile(
         },
       ],
       createImportPlugin(relativeImport),
-      ...(config.babelPluginsPost ?? []),
+      ...[
+        aliases != null
+          ? createAliasRewritePlugin(inputFilePath, aliases)
+          : null,
+        ...(config.babelPluginsPost ?? []),
+      ].filter(Boolean),
     ],
   });
   if (result == null) {