Implements support for python_modules for packages in Python Workers (#9905)

Author: dom96

.changeset/curvy-suns-feel.md

@@ -0,0 +1,5 @@
+---
+"wrangler": minor
+---
+
+Support for Python packages in python_modules dir

packages/wrangler/e2e/dev.test.ts

@@ -515,8 +515,8 @@ describe.each([{ cmd: "wrangler dev" }])(
 					from js import Response
 					def on_fetch(request):
 						return Response.new(f"py hello world {mul(2,3)}")`,
-				"vendor/mod1.py": "print(42)",
-				"vendor/mod2.py": "def hello(): return 42",
+				"python_modules/mod1.py": "print(42)",
+				"python_modules/mod2.py": "def hello(): return 42",
 				"package.json": dedent`
 					{
 						"name": "worker",

packages/wrangler/src/__tests__/config/configuration.test.ts

@@ -29,12 +29,6 @@ describe("readConfig()", () => {
 			    ],
 			    "type": "PythonModule",
 			  },
-			  Object {
-			    "globs": Array [
-			      "vendor/**/*.so",
-			    ],
-			    "type": "Data",
-			  },
 			]
 		`);
 	});

packages/wrangler/src/__tests__/deploy.test.ts

@@ -12007,29 +12007,30 @@ export default{
 
 		it("should print vendor modules correctly in table", async () => {
 			writeWranglerConfig({
-				main: "index.py",
+				main: "src/index.py",
 				compatibility_flags: ["python_workers"],
 			});
 
 			// Create main Python file
 			const mainPython =
 				"from js import Response;\ndef fetch(request):\n return Response.new('hello')";
-			await fs.promises.writeFile("index.py", mainPython);
+			await fs.promises.mkdir("src", { recursive: true });
+			await fs.promises.writeFile("src/index.py", mainPython);
 
 			// Create vendor directory and files
-			await fs.promises.mkdir("vendor", { recursive: true });
+			await fs.promises.mkdir("python_modules", { recursive: true });
 			await fs.promises.writeFile(
-				"vendor/module1.so",
+				"python_modules/module1.so",
 				"binary content for module 1"
 			);
 			await fs.promises.writeFile(
-				"vendor/module2.py",
+				"python_modules/module2.py",
 				"# Python vendor module 2\nprint('hello')"
 			);
 
 			// Create a regular Python module
 			await fs.promises.writeFile(
-				"helper.py",
+				"src/helper.py",
 				"# Helper module\ndef helper(): pass"
 			);
 

packages/wrangler/src/config/validation.ts

@@ -333,16 +333,6 @@ function applyPythonConfig(
 		if (!config.rules.some((rule) => rule.type === "PythonModule")) {
 			config.rules.push({ type: "PythonModule", globs: ["**/*.py"] });
 		}
-		// When vendoring packages they may include certain files that will not be automatically uploaded,
-		// this would require specifying rules in the wrangler configuration of each worker. Instead of
-		// requiring that, we include the config implicitly here.
-		if (
-			!config.rules.some(
-				(rule) => rule.type === "Data" && rule.globs.includes("vendor/**/*.so")
-			)
-		) {
-			config.rules.push({ type: "Data", globs: ["vendor/**/*.so"] });
-		}
 		if (!config.compatibility_flags.includes("python_workers")) {
 			throw new UserError(
 				"The `python_workers` compatibility flag is required to use Python."

packages/wrangler/src/deployment-bundle/find-additional-modules.ts

@@ -57,11 +57,13 @@ function filterPythonVendorModules(
 	if (!isPythonEntrypoint) {
 		return modules;
 	}
-	return modules.filter((m) => !m.name.startsWith("vendor/"));
+	return modules.filter((m) => !m.name.startsWith("python_modules" + path.sep));
 }
 
 function getPythonVendorModulesSize(modules: CfModule[]): number {
-	const vendorModules = modules.filter((m) => m.name.startsWith("vendor/"));
+	const vendorModules = modules.filter((m) =>
+		m.name.startsWith("python_modules" + path.sep)
+	);
 	return vendorModules.reduce((total, m) => total + m.content.length, 0);
 }
 
@@ -136,6 +138,58 @@ export async function findAdditionalModules(
 				filePath: undefined,
 			});
 		}
+
+		// Look for a `python_modules` directory in the root of the project and add all the .py and .so files in it
+		const pythonModulesDir = path.resolve(entry.projectRoot, "python_modules");
+		const pythonModulesDirInModuleRoot = path.resolve(
+			entry.moduleRoot,
+			"python_modules"
+		);
+
+		// Check for conflict between a `python_modules` directory in the module root and the project root.
+		const pythonModulesExistsInModuleRoot = existsSync(
+			pythonModulesDirInModuleRoot
+		);
+		if (
+			pythonModulesExistsInModuleRoot &&
+			entry.projectRoot !== entry.moduleRoot
+		) {
+			throw new UserError(
+				"The 'python_modules' directory cannot exist in your module root. Delete it to continue."
+			);
+		}
+
+		const pythonModulesExists = existsSync(pythonModulesDir);
+		if (pythonModulesExists) {
+			const pythonModulesFiles = getFiles(
+				entry.file,
+				pythonModulesDir,
+				pythonModulesDir,
+				entry.projectRoot
+			);
+			const vendoredRules: Rule[] = [
+				{ type: "Data", globs: ["**/*.so", "**/*.py"] },
+			];
+			const vendoredModules = (
+				await matchFiles(
+					pythonModulesFiles,
+					pythonModulesDir,
+					parseRules(vendoredRules)
+				)
+			).map((m) => {
+				const prefixedPath = path.join("python_modules", m.name);
+				return {
+					...m,
+					name: prefixedPath,
+				};
+			});
+
+			modules.push(...vendoredModules);
+		} else {
+			logger.debug(
+				"Python entrypoint detected, but no python_modules directory found."
+			);
+		}
 	}
 
 	// The modules we find might also have sourcemaps associated with them, so when we go to copy