fix: widen multi-env vars types in wrangler types

Author: dario-piotrowicz

.changeset/olive-kiwis-arrive.md

@@ -0,0 +1,39 @@
+---
+"wrangler": patch
+---
+
+fix: widen multi-env `vars` types in `wrangler types`
+
+Currently types for variable generate string literal, those are appropriate when
+a single environment has been specified in the config file but if multiple environments
+are specified this however wrongly restricts the typing, the changes here fix such
+incorrect behavior.
+
+For example, given a `wrangler.toml` containing the following:
+
+```
+[vars]
+MY_VAR = "dev value"
+
+[env.production]
+[env.production.vars]
+MY_VAR = "prod value"
+```
+
+running `wrangler types` would generate:
+
+```ts
+interface Env {
+	MY_VAR: "dev value";
+}
+```
+
+making typescript incorrectly assume that `MY_VAR` is always going to be `"dev value"`
+
+after these changes, the generated interface would instead be:
+
+```ts
+interface Env {
+	MY_VAR: "dev value" | "prod value";
+}
+```

packages/wrangler/src/__tests__/type-generation.test.ts

@@ -391,6 +391,48 @@ describe("generateTypes()", () => {
 	`);
 	});
 
+	it("should produce unions where appropriate for vars present in multiple environments", async () => {
+		fs.writeFileSync(
+			"./wrangler.toml",
+			TOML.stringify({
+				vars: {
+					MY_VAR: "a var",
+					MY_VAR_A: "A (dev)",
+					MY_VAR_B: { value: "B (dev)" },
+					MY_VAR_C: ["a", "b", "c"],
+				},
+				env: {
+					production: {
+						vars: {
+							MY_VAR: "a var",
+							MY_VAR_A: "A (prod)",
+							MY_VAR_B: { value: "B (prod)" },
+							MY_VAR_C: [1, 2, 3],
+						},
+					},
+					staging: {
+						vars: {
+							MY_VAR_A: "A (stag)",
+						},
+					},
+				},
+			} as TOML.JsonMap),
+			"utf-8"
+		);
+
+		await runWrangler("types");
+
+		expect(std.out).toMatchInlineSnapshot(`
+		"interface Env {
+			MY_VAR: \\"a var\\";
+			MY_VAR_A: \\"A (dev)\\" | \\"A (prod)\\" | \\"A (stag)\\";
+			MY_VAR_C: [\\"a\\",\\"b\\",\\"c\\"] | [1,2,3];
+			MY_VAR_B: {\\"value\\":\\"B (dev)\\"} | {\\"value\\":\\"B (prod)\\"};
+		}
+		"
+	`);
+	});
+
 	describe("customization", () => {
 		describe("env", () => {
 			it("should allow the user to customize the interface name", async () => {

packages/wrangler/src/config/index.ts

@@ -31,16 +31,11 @@ export function readConfig<CommandArgs>(
 	args: CommandArgs &
 		Pick<OnlyCamelCase<CommonYargsOptions>, "experimentalJsonConfig">
 ): Config {
-	let rawConfig: RawConfig = {};
-	if (!configPath) {
-		configPath = findWranglerToml(process.cwd(), args.experimentalJsonConfig);
-	}
-	// Load the configuration from disk if available
-	if (configPath?.endsWith("toml")) {
-		rawConfig = parseTOML(readFileSync(configPath), configPath);
-	} else if (configPath?.endsWith("json")) {
-		rawConfig = parseJSONC(readFileSync(configPath), configPath);
-	}
+	const { rawConfig, configPath: updateConfigPath } = getRawConfig<CommandArgs>(
+		configPath,
+		args
+	);
+	configPath = updateConfigPath;
 
 	// Process the top-level configuration.
 	const { config, diagnostics } = normalizeAndValidateConfig(
@@ -70,6 +65,30 @@ export function readConfig<CommandArgs>(
 	return config;
 }
 
+/**
+ * Get the raw (pre-processed) Wrangler configuration; read it from the give `configPath` if available.
+ * Returns the raw configuration alongside the `configPath` for it (which is equal to the provided one, if one was)
+ */
+export function getRawConfig<CommandArgs>(
+	configPath: string | undefined,
+	args: CommandArgs &
+		Pick<OnlyCamelCase<CommonYargsOptions>, "experimentalJsonConfig">
+) {
+	let rawConfig: RawConfig = {};
+	if (!configPath) {
+		configPath = findWranglerToml(process.cwd(), args.experimentalJsonConfig);
+	}
+
+	// Load the configuration from disk if available
+	if (configPath?.endsWith("toml")) {
+		rawConfig = parseTOML(readFileSync(configPath), configPath);
+	} else if (configPath?.endsWith("json")) {
+		rawConfig = parseJSONC(readFileSync(configPath), configPath);
+	}
+
+	return { rawConfig, configPath };
+}
+
 /**
  * Find the wrangler.toml file by searching up the file-system
  * from the current working directory.

packages/wrangler/src/type-generation.ts

@@ -1,13 +1,13 @@
 import * as fs from "node:fs";
 import { findUpSync } from "find-up";
-import { findWranglerToml, readConfig } from "./config";
+import { findWranglerToml, getRawConfig, readConfig } from "./config";
 import { getEntry } from "./deployment-bundle/entry";
 import { getVarsForDev } from "./dev/dev-vars";
 import { UserError } from "./errors";
 import { logger } from "./logger";
 import { printWranglerBanner } from "./update-check";
 import { CommandLineArgsError } from "./index";
-import type { Config } from "./config";
+import type { Config, RawEnvironment } from "./config";
 import type { CfScriptFormat } from "./deployment-bundle/worker";
 import type {
 	CommonYargsArgv,
@@ -74,11 +74,9 @@ export async function typesHandler(
 		true
 	) as Record<string, string>;
 
-	const configBindingsWithSecrets: Partial<Config> & {
-		secrets: Record<string, string>;
-	} = {
+	const configBindingsWithSecrets: ConfigToDTS = {
 		kv_namespaces: config.kv_namespaces ?? [],
-		vars: { ...config.vars },
+		vars: getVarsInfo(configPath, args),
 		wasm_modules: config.wasm_modules,
 		text_blobs: {
 			...config.text_blobs,
@@ -108,8 +106,12 @@ export async function typesHandler(
 
 type Secrets = Record<string, string>;
 
+type ConfigToDTS = Partial<Omit<Config, "vars">> & { vars: VarsInfo } & {
+	secrets: Secrets;
+};
+
 async function generateTypes(
-	configToDTS: Partial<Config> & { secrets: Secrets },
+	configToDTS: ConfigToDTS,
 	config: Config,
 	envInterface: string,
 	outputPath: string
@@ -148,17 +150,25 @@ async function generateTypes(
 		const vars = Object.entries(configToDTS.vars).filter(
 			([key]) => !(key in configToDTS.secrets)
 		);
-		for (const [varName, varValue] of vars) {
-			if (
-				typeof varValue === "string" ||
-				typeof varValue === "number" ||
-				typeof varValue === "boolean"
-			) {
-				envTypeStructure.push(`${varName}: "${varValue}";`);
-			}
-			if (typeof varValue === "object" && varValue !== null) {
-				envTypeStructure.push(`${varName}: ${JSON.stringify(varValue)};`);
-			}
+		for (const [varName, varInfo] of vars) {
+			const varValueTypes = new Set(
+				varInfo
+					.map(({ value }) => value)
+					.map((varValue) => {
+						if (
+							typeof varValue === "string" ||
+							typeof varValue === "number" ||
+							typeof varValue === "boolean"
+						) {
+							return `"${varValue}"`;
+						}
+						if (typeof varValue === "object" && varValue !== null) {
+							return `${JSON.stringify(varValue)}`;
+						}
+					})
+					.filter(Boolean)
+			);
+			envTypeStructure.push(`${varName}: ${[...varValueTypes].join(" | ")};`);
 		}
 	}
 
@@ -329,3 +339,31 @@ function writeDTSFile({
 		logger.log(combinedTypeStrings);
 	}
 }
+
+type VarValue = Config["vars"][string];
+
+type VarInfoValue = { value: VarValue; env?: string };
+
+type VarsInfo = Record<string, VarInfoValue[]>;
+
+function getVarsInfo(
+	configPath: string,
+	args: StrictYargsOptionsToInterface<typeof typesOptions>
+): VarsInfo {
+	const varsInfo: VarsInfo = {};
+	const { rawConfig } = getRawConfig(configPath, args);
+
+	function collectVars(vars: RawEnvironment["vars"], envName?: string) {
+		Object.entries(vars ?? {}).forEach(([key, value]) => {
+			varsInfo[key] ??= [];
+			varsInfo[key].push({ value, env: envName });
+		});
+	}
+
+	collectVars(rawConfig.vars);
+	Object.entries(rawConfig.env ?? {}).forEach(([envName, env]) => {
+		collectVars(env.vars, envName);
+	});
+
+	return varsInfo;
+}