Add add_ons metadata to pyproject.toml for project creation (#3188)

* release note Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * refactor out the add_ons_dict constant Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * refactor the parsing logic and move to cli/starters.py from tepmplate Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * use the name of the add_ons Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * remove parsing steps Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * add comments Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * fix typo Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * fix type hint Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * fix release note Signed-off-by: Nok <nok.lam.chan@quantumblack.com> * adhoc fix Signed-off-by: Nok <nok.lam.chan@quantumblack.com> --------- Signed-off-by: Nok <nok.lam.chan@quantumblack.com> Signed-off-by: Sajid Alam <90610031+SajidAlamQB@users.noreply.github.com> Signed-off-by: Nok Lam Chan <nok.lam.chan@quantumblack.com> Co-authored-by: Sajid Alam <90610031+SajidAlamQB@users.noreply.github.com>
kedro-org · Oct 24, 2023 · 00f035a · 00f035a
1 parent 05f1473
commit 00f035a
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 82 deletions.
diff --git a/RELEASE.md b/RELEASE.md
@@ -6,6 +6,7 @@
 * The new spaceflights starters, `spaceflights-pandas`, `spaceflights-pandas-viz`, `spaceflights-pyspark`, and `spaceflights-pyspark-viz` can be used with the `kedro new` command with the `--starter` flag.
 
 ## Bug fixes and other changes
+* Added a new field `add-ons` to `pyproject.toml` when a project is created.
 
 ## Breaking changes to the API
 * Renamed the `data_sets` argument and the `_data_sets` attribute in `Catalog` and their references to `datasets` and `_datasets` respectively.

diff --git a/kedro/framework/cli/starters.py b/kedro/framework/cli/starters.py
@@ -32,7 +32,6 @@
     _safe_load_entry_point,
     command_with_verbosity,
 )
-from kedro.templates.project.hooks.utils import parse_add_ons_input
 
 KEDRO_PATH = Path(kedro.__file__).parent
 TEMPLATE_PATH = KEDRO_PATH / "templates" / "project"
@@ -108,7 +107,13 @@ class KedroStarterSpec:  # noqa: too-few-public-methods
     "An optional directory inside the repository where the starter resides."
 )
 
-
+ADD_ONS_DICT = {
+    "1": "Linting",
+    "2": "Testing",
+    "3": "Custom Logging",
+    "4": "Documentation",
+    "5": "Data Structure",
+}
 # noqa: unused-argument
 def _remove_readonly(func: Callable, path: Path, excinfo: tuple):  # pragma: no cover
     """Remove readonly files on Windows
@@ -181,6 +186,50 @@ def _starter_spec_to_dict(
     return format_dict
 
 
+def _parse_add_ons_input(add_ons_str: str):
+    """Parse the add-ons input string.
+
+    Args:
+        add_ons_str: Input string from prompts.yml.
+
+    Returns:
+        list: List of selected add-ons as strings.
+    """
+
+    def _validate_range(start, end):
+        if int(start) > int(end):
+            message = f"'{start}-{end}' is an invalid range for project add-ons.\nPlease ensure range values go from smaller to larger."
+            click.secho(message, fg="red", err=True)
+            sys.exit(1)
+
+    def _validate_selection(add_ons: list[str]):
+        for add_on in add_ons:
+            if int(add_on) < 1 or int(add_on) > len(ADD_ONS_DICT):
+                message = f"'{add_on}' is not a valid selection.\nPlease select from the available add-ons: 1, 2, 3, 4, 5."  # nosec
+                click.secho(message, fg="red", err=True)
+                sys.exit(1)
+
+    if add_ons_str == "all":
+        return list(ADD_ONS_DICT)
+    if add_ons_str == "none":
+        return []
+
+    # Split by comma
+    add_ons_choices = add_ons_str.split(",")
+    selected: list[str] = []
+
+    for choice in add_ons_choices:
+        if "-" in choice:
+            start, end = choice.split("-")
+            _validate_range(start, end)
+            selected.extend(str(i) for i in range(int(start), int(end) + 1))
+        else:
+            selected.append(choice.strip())
+
+    _validate_selection(selected)
+    return selected
+
+
 # noqa: missing-function-docstring
 @click.group(context_settings=CONTEXT_SETTINGS, name="Kedro")
 def create_cli():  # pragma: no cover
@@ -337,7 +386,7 @@ def _fetch_config_from_file(config_path: str) -> dict[str, str]:
 
 
 def _make_cookiecutter_args(
-    config: dict[str, str],
+    config: dict[str, str | list[str]],
     checkout: str,
     directory: str,
 ) -> dict[str, Any]:
@@ -360,11 +409,20 @@ def _make_cookiecutter_args(
     """
     config.setdefault("kedro_version", version)
 
+    # Map the selected add on lists to readable name
+    add_ons = config.get("add_ons")
+    if add_ons:
+        config["add_ons"] = [
+            ADD_ONS_DICT[add_on] for add_on in _parse_add_ons_input(add_ons)  # type: ignore
+        ]
+        config["add_ons"] = str(config["add_ons"])
+
     cookiecutter_args = {
         "output_dir": config.get("output_dir", str(Path.cwd().resolve())),
         "no_input": True,
         "extra_context": config,
     }
+
     if checkout:
         cookiecutter_args["checkout"] = checkout
     if directory:
@@ -373,22 +431,6 @@ def _make_cookiecutter_args(
     return cookiecutter_args
 
 
-def _get_add_ons_text(add_ons):
-    add_ons_dict = {
-        "1": "Linting",
-        "2": "Testing",
-        "3": "Custom Logging",
-        "4": "Documentation",
-        "5": "Data structure",
-    }
-    add_ons_list = parse_add_ons_input(add_ons)
-    add_ons_text = [add_ons_dict[add_on] for add_on in add_ons_list]
-    return (
-        " ".join(str(add_on) + "," for add_on in add_ons_text[:-1])
-        + f" and {add_ons_text[-1]}"
-    )
-
-
 def _create_project(template_path: str, cookiecutter_args: dict[str, Any]):
     """Creates a new kedro project using cookiecutter.
 
@@ -422,12 +464,10 @@ def _create_project(template_path: str, cookiecutter_args: dict[str, Any]):
 
     # Only non-starter projects have configurable add-ons
     if template_path == str(TEMPLATE_PATH):
-        if add_ons == "none":
+        if add_ons == "[]":  # TODO: This should be a list
             click.secho("\nYou have selected no add-ons")
         else:
-            click.secho(
-                f"\nYou have selected the following add-ons: {_get_add_ons_text(add_ons)}"
-            )
+            click.secho(f"\nYou have selected the following add-ons: {add_ons}")
 
     click.secho(
         f"\nThe project name '{project_name}' has been applied to: "

diff --git a/kedro/templates/project/hooks/post_gen_project.py b/kedro/templates/project/hooks/post_gen_project.py
@@ -1,10 +1,11 @@
 from pathlib import Path
 
 from kedro.templates.project.hooks.utils import (
-    parse_add_ons_input,
+
     setup_template_add_ons,
     sort_requirements,
 )
+from kedro.framework.cli.starters import _parse_add_ons_input
 
 def main():
     current_dir = Path.cwd()
@@ -14,11 +15,8 @@ def main():
     # Get the selected add-ons from cookiecutter
     selected_add_ons = "{{ cookiecutter.add_ons }}"
 
-    # Parse the add-ons to get a list
-    selected_add_ons_list = parse_add_ons_input(selected_add_ons)
-
     # Handle template directories and requirements according to selected add-ons
-    setup_template_add_ons(selected_add_ons_list, requirements_file_path, pyproject_file_path)
+    setup_template_add_ons(selected_add_ons, requirements_file_path, pyproject_file_path)
 
     # Sort requirements.txt file in alphabetical order
     sort_requirements(requirements_file_path)

diff --git a/kedro/templates/project/hooks/utils.py b/kedro/templates/project/hooks/utils.py
@@ -47,49 +47,6 @@
 ]
 """
 
-def _validate_range(start, end):
-    if int(start) > int(end):
-        message = f"'{start}-{end}' is an invalid range for project add-ons.\nPlease ensure range values go from smaller to larger."
-        click.secho(message, fg="red", err=True)
-        sys.exit(1)
-
-def _validate_selection(add_ons):
-    for add_on in add_ons:
-        if int(add_on) < 1 or int(add_on) > 5:
-            message = f"'{add_on}' is not a valid selection.\nPlease select from the available add-ons: 1, 2, 3, 4, 5."
-            click.secho(message, fg="red", err=True)
-            sys.exit(1)
-
-
-def parse_add_ons_input(add_ons_str):
-    """Parse the add-ons input string.
-
-    Args:
-        add_ons_str: Input string from prompts.yml.
-
-    Returns:
-        list: List of selected add-ons as strings.
-    """
-    if add_ons_str == "all":
-        return ["1", "2", "3", "4", "5"]
-    if add_ons_str == "none":
-        return []
-
-    # Split by comma
-    add_ons_choices = add_ons_str.split(",")
-    selected = []
-
-    for choice in add_ons_choices:
-        if "-" in choice:
-            start, end = choice.split("-")
-            _validate_range(start, end)
-            selected.extend(str(i) for i in range(int(start), int(end) + 1))
-        else:
-            selected.append(choice.strip())
-
-    _validate_selection(selected)
-    return selected
-
 
 def setup_template_add_ons(selected_add_ons_list, requirements_file_path, pyproject_file_path):
     """Removes directories and files related to unwanted addons from
@@ -103,15 +60,15 @@ def setup_template_add_ons(selected_add_ons_list, requirements_file_path, pyproj
         pyproject_file_path: the path to the pyproject.toml file
             located on the the root of the template.
     """
-    if "1" not in selected_add_ons_list:  # If Linting not selected
+    if "Linting" not in selected_add_ons_list:
         pass
     else:
         with open(requirements_file_path, 'a') as file:
             file.write(lint_requirements)
         with open(pyproject_file_path, 'a') as file:
             file.write(lint_pyproject_requirements)
 
-    if "2" not in selected_add_ons_list:  # If Testing not selected
+    if "Testing" not in selected_add_ons_list:
         tests_path = current_dir / "tests"
         if tests_path.exists():
             shutil.rmtree(str(tests_path))
@@ -121,20 +78,20 @@ def setup_template_add_ons(selected_add_ons_list, requirements_file_path, pyproj
         with open(pyproject_file_path, 'a') as file:
             file.write(test_pyproject_requirements)
 
-    if "3" not in selected_add_ons_list:  # If Logging not selected
+    if "Logging" not in selected_add_ons_list:
         logging_yml_path = current_dir / "conf/logging.yml"
         if logging_yml_path.exists():
             logging_yml_path.unlink()
 
-    if "4" not in selected_add_ons_list:  # If Documentation not selected
+    if "Documentation" not in selected_add_ons_list:
         docs_path = current_dir / "docs"
         if docs_path.exists():
             shutil.rmtree(str(docs_path))
     else:
         with open(pyproject_file_path, 'a') as file:
             file.write(docs_pyproject_requirements)
 
-    if "5" not in selected_add_ons_list:  # If Data Structure not selected
+    if "Data Structure" not in selected_add_ons_list:
         data_path = current_dir / "data"
         if data_path.exists():
             shutil.rmtree(str(data_path))

diff --git a/kedro/templates/project/{{ cookiecutter.repo_name }}/pyproject.toml b/kedro/templates/project/{{ cookiecutter.repo_name }}/pyproject.toml
@@ -37,3 +37,4 @@ namespaces = false
 package_name = "{{ cookiecutter.python_package }}"
 project_name = "{{ cookiecutter.project_name }}"
 kedro_init_version = "{{ cookiecutter.kedro_version }}"
+add_ons = {{ cookiecutter.add_ons | string | replace("\'", "\"") }}
diff --git a/tests/framework/cli/test_starters.py b/tests/framework/cli/test_starters.py
@@ -15,8 +15,8 @@
     _OFFICIAL_STARTER_SPECS,
     TEMPLATE_PATH,
     KedroStarterSpec,
+    _parse_add_ons_input,
 )
-from kedro.templates.project.hooks.utils import parse_add_ons_input
 
 FILES_IN_TEMPLATE_WITH_NO_ADD_ONS = 14
 
@@ -64,7 +64,7 @@ def _get_expected_files(add_ons: str):
         "4": 2,
         "5": 8,
     }  # files added to template by each add-on
-    add_ons_list = parse_add_ons_input(add_ons)
+    add_ons_list = _parse_add_ons_input(add_ons)
 
     expected_files = FILES_IN_TEMPLATE_WITH_NO_ADD_ONS
 
@@ -88,7 +88,7 @@ def _assert_requirements_ok(
     requirements_file_path = root_path / "requirements.txt"
     pyproject_file_path = root_path / "pyproject.toml"
 
-    add_ons_list = parse_add_ons_input(add_ons)
+    add_ons_list = _parse_add_ons_input(add_ons)
 
     if "1" in add_ons_list:
         with open(requirements_file_path) as requirements_file:
@@ -256,7 +256,7 @@ def test_starter_list_with_invalid_starter_plugin(
     ],
 )
 def test_parse_add_ons_valid(input, expected):
-    result = parse_add_ons_input(input)
+    result = _parse_add_ons_input(input)
     assert result == expected
 
 
@@ -266,7 +266,7 @@ def test_parse_add_ons_valid(input, expected):
 )
 def test_parse_add_ons_invalid_range(input, capsys):
     with pytest.raises(SystemExit):
-        parse_add_ons_input(input)
+        _parse_add_ons_input(input)
     message = f"'{input}' is an invalid range for project add-ons.\nPlease ensure range values go from smaller to larger."
     assert message in capsys.readouterr().err
 
@@ -277,7 +277,7 @@ def test_parse_add_ons_invalid_range(input, capsys):
 )
 def test_parse_add_ons_invalid_selection(input, first_invalid, capsys):
     with pytest.raises(SystemExit):
-        parse_add_ons_input(input)
+        _parse_add_ons_input(input)
     message = f"'{first_invalid}' is not a valid selection.\nPlease select from the available add-ons: 1, 2, 3, 4, 5."
     assert message in capsys.readouterr().err