From 422730d51600178d49867975b59d1d3541a03c09 Mon Sep 17 00:00:00 2001 From: konsti Date: Tue, 20 Aug 2024 14:41:07 +0200 Subject: [PATCH] Benchmarks: Add extras to poetry (#6247) Previously, we had dropped extras from the requirements when converting to poetry, skewing the results towards poetry for the `apache-airflow[all]` benchmark. --- scripts/benchmark/src/benchmark/resolver.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/scripts/benchmark/src/benchmark/resolver.py b/scripts/benchmark/src/benchmark/resolver.py index 85cfc988b..923a4c313 100644 --- a/scripts/benchmark/src/benchmark/resolver.py +++ b/scripts/benchmark/src/benchmark/resolver.py @@ -374,14 +374,13 @@ class Poetry(Suite): pyproject = tomli.load(fp) # Add the dependencies to the pyproject.toml. - pyproject["tool"]["poetry"]["dependencies"].update( - { - str(requirement.name): str(requirement.specifier) - if requirement.specifier - else "*" - for requirement in requirements - } - ) + for requirement in requirements: + version = str(requirement.specifier) if requirement.specifier else "*" + if requirement.extras: + entry = {"version": version, "extras": sorted(requirement.extras)} + else: + entry = version + pyproject["tool"]["poetry"]["dependencies"][requirement.name] = entry with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: tomli_w.dump(pyproject, fp)