From 76ddd66264b753402ad4ddcafb400f9a7a272996 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikael=20G=C3=B6ransson?= Date: Thu, 1 Feb 2024 13:48:16 +0100 Subject: [PATCH 1/2] use `importlib.util` to import module from locustfile instead of `load_module`, which is deprecated and was planned to be removed in 3.12. --- locust/util/load_locustfile.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/locust/util/load_locustfile.py b/locust/util/load_locustfile.py index db3b63c037..423ee47380 100644 --- a/locust/util/load_locustfile.py +++ b/locust/util/load_locustfile.py @@ -1,6 +1,7 @@ from __future__ import annotations import importlib +import importlib.util import inspect import os import sys @@ -53,9 +54,15 @@ def load_locustfile(path) -> tuple[str | None, dict[str, User], list[LoadTestSha # Add to front, then remove from original position sys.path.insert(0, directory) del sys.path[i + 1] + # Perform the import - source = importlib.machinery.SourceFileLoader(os.path.splitext(locustfile)[0], path) - imported = source.load_module() + module_name = os.path.splitext(locustfile)[0] + loader = importlib.machinery.SourceFileLoader(module_name, path) + spec = importlib.util.spec_from_file_location(module_name, path, loader=loader) + imported = importlib.util.module_from_spec(spec) + sys.modules[imported.__name__] = imported + loader.exec_module(imported) + # Remove directory from path if we added it ourselves (just to be neat) if added_to_path: del sys.path[0] @@ -67,6 +74,6 @@ def load_locustfile(path) -> tuple[str | None, dict[str, User], list[LoadTestSha user_classes = {name: value for name, value in vars(imported).items() if is_user_class(value)} # Find shape class, if any, return it - shape_classes = [value() for name, value in vars(imported).items() if is_shape_class(value)] + shape_classes = [value() for value in vars(imported).values() if is_shape_class(value)] return imported.__doc__, user_classes, shape_classes From 1d768c814c882fc42f67062844400794947acbcf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikael=20G=C3=B6ransson?= Date: Thu, 1 Feb 2024 13:55:07 +0100 Subject: [PATCH 2/2] handle if no module spec was found --- locust/util/load_locustfile.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/locust/util/load_locustfile.py b/locust/util/load_locustfile.py index 423ee47380..3c9904945c 100644 --- a/locust/util/load_locustfile.py +++ b/locust/util/load_locustfile.py @@ -59,6 +59,10 @@ def load_locustfile(path) -> tuple[str | None, dict[str, User], list[LoadTestSha module_name = os.path.splitext(locustfile)[0] loader = importlib.machinery.SourceFileLoader(module_name, path) spec = importlib.util.spec_from_file_location(module_name, path, loader=loader) + if spec is None: + sys.stderr.write(f"Unable to get module spec for {module_name} in {path}") + sys.exit(1) + imported = importlib.util.module_from_spec(spec) sys.modules[imported.__name__] = imported loader.exec_module(imported)