From 3fad7aee4503793a0645d075f58dd6533be6fd21 Mon Sep 17 00:00:00 2001
From: Morgan Chorlton <morganchorlton3@gmail.com>
Date: Tue, 11 Jun 2024 14:49:26 +0100
Subject: [PATCH] Updated argument to accept additional poetry arguments

---
 package.py | 49 +++++++++++++++++--------------------------------
 1 file changed, 17 insertions(+), 32 deletions(-)

diff --git a/package.py b/package.py
index 028c970a..10537030 100644
--- a/package.py
+++ b/package.py
@@ -693,7 +693,7 @@ def pip_requirements_step(path, prefix=None, required=False, tmp_dir=None):
                 step("pip", runtime, requirements, prefix, tmp_dir)
                 hash(requirements)
 
-        def poetry_install_step(path, poetry_groups=None, prefix=None, required=False):
+        def poetry_install_step(path, poetry_args=[], prefix=None, required=False):
             pyproject_file = path
             if os.path.isdir(path):
                 pyproject_file = os.path.join(path, "pyproject.toml")
@@ -703,7 +703,7 @@ def poetry_install_step(path, poetry_groups=None, prefix=None, required=False):
                         "poetry configuration not found: {}".format(pyproject_file)
                     )
             else:
-                step("poetry", runtime, path, poetry_groups, prefix)
+                step("poetry", runtime, path, poetry_args, prefix)
                 hash(pyproject_file)
                 pyproject_path = os.path.dirname(pyproject_file)
                 poetry_lock_file = os.path.join(pyproject_path, "poetry.lock")
@@ -807,7 +807,7 @@ def commands_step(path, commands):
                     prefix = claim.get("prefix_in_zip")
                     pip_requirements = claim.get("pip_requirements")
                     poetry_install = claim.get("poetry_install")
-                    poetry_groups = claim.get("poetry_groups")
+                    additional_poetry_arg = claim.get("additional_poetry_args", [])
                     npm_requirements = claim.get("npm_package_json")
                     runtime = claim.get("runtime", query.runtime)
 
@@ -832,7 +832,7 @@ def commands_step(path, commands):
                             poetry_install_step(
                                 path,
                                 prefix=prefix,
-                                poetry_groups=poetry_groups,
+                                poetry_args=additional_poetry_arg,
                                 required=True,
                             )
 
@@ -907,11 +907,11 @@ def execute(self, build_plan, zip_stream, query):
                 (
                     runtime,
                     path,
-                    poetry_groups,
+                    poetry_args,
                     prefix,
                 ) = action[1:]
-                log.info("Poetry Groups: %s", poetry_groups)
-                with install_poetry_dependencies(query, path, poetry_groups) as rd:
+                log.info("Poetry arguments: %s", poetry_args)
+                with install_poetry_dependencies(query, path, poetry_args) as rd:
                     if rd:
                         if pf:
                             self._zip_write_with_filter(zs, pf, rd, prefix, timestamp=0)
@@ -1106,7 +1106,7 @@ def install_pip_requirements(query, requirements_file, tmp_dir):
 
 
 @contextmanager
-def install_poetry_dependencies(query, path, poetry_groups):
+def install_poetry_dependencies(query, path, poetry_args):
     # TODO:
     #  1. Emit files instead of temp_dir
 
@@ -1195,31 +1195,16 @@ def copy_file_to_target(file, temp_dir):
             # NOTE: poetry must be available in the build environment, which is the case with lambci/lambda:build-python* docker images but not public.ecr.aws/sam/build-python* docker images
             # FIXME: poetry install does not currently allow to specify the target directory so we export the
             # requirements then install them with "pip --no-deps" to avoid using pip dependency resolver
-            if poetry_groups is not None:
-                group_args = []
-                for group in poetry_groups:
-                    group_args.append("--with")
-                    group_args.append(group)
 
-                poetry_export = [
-                    poetry_exec,
-                    "export",
-                    "--format",
-                    "requirements.txt",
-                    "--output",
-                    "requirements.txt",
-                    "--with-credentials",
-                ] + group_args
-            else:
-                poetry_export = [
-                    poetry_exec,
-                    "export",
-                    "--format",
-                    "requirements.txt",
-                    "--output",
-                    "requirements.txt",
-                    "--with-credentials",
-                ]
+            poetry_export = [
+                poetry_exec,
+                "export",
+                "--format",
+                "requirements.txt",
+                "--output",
+                "requirements.txt",
+                "--with-credentials",
+            ] + poetry_args
 
             poetry_commands = [
                 [