📦 colinhacks / zshy

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144import * as fs from "node:fs";
import * as path from "node:path";
import * as ts from "typescript";
import * as utils from "./utils.js";

export const createExtensionRewriteTransformer =
  (config: {
    rootDir: string;
    ext: string;
    onAssetImport?: (assetPath: string) => void;
  }): ts.TransformerFactory<ts.SourceFile | ts.Bundle> =>
  (context) => {
    return (sourceFile) => {
      const visitor = (node: ts.Node): ts.Node => {
        const isImport = ts.isImportDeclaration(node);
        const isExport = ts.isExportDeclaration(node);
        const isDynamicImport = ts.isCallExpression(node) && node.expression.kind === ts.SyntaxKind.ImportKeyword;

        let originalText: string;
        if (isImport || isExport || isDynamicImport) {
          if (isImport || isExport) {
            if (!node.moduleSpecifier || !ts.isStringLiteral(node.moduleSpecifier)) {
              return ts.visitEachChild(node, visitor, context);
            }

            originalText = node.moduleSpecifier.text;
          } else if (isDynamicImport) {
            const arg = node.arguments[0]!;
            if (!ts.isStringLiteral(arg)) {
              // continue
              return ts.visitEachChild(node, visitor, context);
            }
            originalText = arg.text;
          } else {
            // If it's not an import, export, or dynamic import, just visit children
            return ts.visitEachChild(node, visitor, context);
          }

          const isRelativeImport = originalText.startsWith("./") || originalText.startsWith("../");
          if (!isRelativeImport) {
            // If it's not a relative import, don't transform it
            return node;
          }

          const ext = path.extname(originalText).toLowerCase();

          // rewrite .js to resolved js extension
          if (ext === ".js" || ext === ".ts") {
            const newText = originalText.slice(0, -3) + config.ext;
            if (isImport) {
              return ts.factory.updateImportDeclaration(
                node,
                node.modifiers,
                node.importClause,
                ts.factory.createStringLiteral(newText),
                node.assertClause
              );
            } else if (isExport) {
              return ts.factory.updateExportDeclaration(
                node,
                node.modifiers,
                node.isTypeOnly,
                node.exportClause,
                ts.factory.createStringLiteral(newText),
                node.assertClause
              );
            } else if (isDynamicImport) {
              return ts.factory.updateCallExpression(node, node.expression, node.typeArguments, [
                ts.factory.createStringLiteral(newText),
                ...node.arguments.slice(1),
              ]);
            }
          }

          // rewrite extensionless imports to .js
          if (ext === "") {
            // Check filesystem to determine if we should resolve to file.ts or directory/index.ts
            let newText = originalText + config.ext;

            if (ts.isSourceFile(sourceFile)) {
              const sourceFileDir = path.dirname(sourceFile.fileName);
              const resolvedPath = path.resolve(sourceFileDir, originalText);

              // Check if the extensionless import refers to a file (e.g., d.ts)
              const potentialFile = resolvedPath + ".ts";
              const potentialIndexFile = path.join(resolvedPath, "index.ts");

              if (fs.existsSync(potentialIndexFile) && !fs.existsSync(potentialFile)) {
                // Directory with index.ts exists, use index path
                newText = originalText + "/index" + config.ext;
              }
              // Otherwise, use the default behavior (originalText + config.ext)
            }

            if (isImport) {
              return ts.factory.updateImportDeclaration(
                node,
                node.modifiers,
                node.importClause,
                ts.factory.createStringLiteral(newText),
                node.assertClause
              );
            } else if (isExport) {
              return ts.factory.updateExportDeclaration(
                node,
                node.modifiers,
                node.isTypeOnly,
                node.exportClause,
                ts.factory.createStringLiteral(newText),
                node.assertClause
              );
            } else if (isDynamicImport) {
              return ts.factory.updateCallExpression(node, node.expression, node.typeArguments, [
                ts.factory.createStringLiteral(newText),
                ...node.arguments.slice(1),
              ]);
            }
          }

          // copy asset files
          if (utils.isAssetFile(originalText)) {
            // it's an asset
            if (ts.isSourceFile(sourceFile)) {
              const sourceFileDir = path.dirname(sourceFile.fileName);
              const resolvedAssetPath = path.resolve(sourceFileDir, originalText);
              // Make it relative to the source root (rootDir)
              const relAssetPath = path.relative(config.rootDir, resolvedAssetPath);
              // Track asset import if callback provided
              if (config.onAssetImport) {
                config.onAssetImport(relAssetPath);
              }
            }
            // Don't transform asset dynamic imports, leave them as-is
            return node;
          }
        }

        return ts.visitEachChild(node, visitor, context);
      };

      return ts.visitNode(sourceFile, visitor) as ts.SourceFile;
    };
  };