Updates to pyflyte init templating

flytekit/clis/sdk_in_container/helpers.py

@@ -1,7 +1,10 @@
+import os
+import shutil
 from dataclasses import replace
 from typing import Optional
 
 import rich_click as click
+from git import Repo
 
 from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE
 from flytekit.configuration import Config, ImageConfig, get_config_file
@@ -63,3 +66,34 @@ def patch_image_config(config_file: Optional[str], image_config: ImageConfig) ->
             if addl.name not in additional_image_names:
                 new_additional_images.append(addl)
     return replace(image_config, default_image=new_default, images=new_additional_images)
+
+
+def clone_and_copy_repo_dir(git_url, branch, src_dir, dest_dir):
+    """
+    Clones a git repository, checks out to a specific branch, and copies a specified directory into a local directory.
+
+    Parameters:
+    - github_url: The URL of the GitHub repository.
+    - branch: The branch of the repository.
+    - src_dir: The directory in the repository to copy.
+    - dest_dir: The local directory to copy into.
+    """
+    # Clone the repo
+    repo = Repo.clone_from(git_url, "temp_repo")
+
+    # Checkout to the branch
+    repo.git.checkout(branch)
+
+    # Define the source directory path
+    src_path = os.path.join("temp_repo", src_dir)
+
+    # Check if source directory exists
+    if not os.path.exists(src_path):
+        print(f"Template named {src_dir} does not exist in the repository.")
+        return
+
+    # Copy the files from the source directory to the destination directory
+    shutil.copytree(src_path, dest_dir)
+
+    # Remove the temporary cloned repo
+    shutil.rmtree("temp_repo")

flytekit/clis/sdk_in_container/init.py

@@ -1,5 +1,6 @@
 import rich_click as click
-from cookiecutter.main import cookiecutter
+
+from flytekit.clis.sdk_in_container.helpers import clone_and_copy_repo_dir
 
 
 @click.command("init")
@@ -8,29 +9,19 @@
     default="simple-example",
     help="cookiecutter template folder name to be used in the repo - https://github.com/flyteorg/flytekit-python-template.git",
 )
+@click.option(
+    "--repository-url",
+    default="https://github.com/flyteorg/flytekit-python-template.git",
+    help="template repository url pointing to a git repository containing flytekit templates.",
+)
+@click.option("--repository-branch", default="main", help="template repository branch to be used.")
 @click.argument("project-name")
-def init(template, project_name):
+def init(template, repository_url, repository_branch, project_name):
     """
     Create flyte-ready projects.
     """
-    config = {
-        "project_name": project_name,
-        "app": "flyte",
-        "workflow": "my_wf",
-    }
-    cookiecutter(
-        "https://github.com/flyteorg/flytekit-python-template.git",
-        checkout="main",
-        no_input=True,
-        # We do not want to clobber existing files/directories.
-        overwrite_if_exists=False,
-        extra_context=config,
-        # By specifying directory we can have multiple templates in the same repository,
-        # as described in https://cookiecutter.readthedocs.io/en/1.7.2/advanced/directories.html.
-        # The idea is to extend the number of templates, each in their own subdirectory, for example
-        # a tensorflow-based example.
-        directory=template,
-    )
+
+    clone_and_copy_repo_dir(repository_url, repository_branch, template, project_name)
 
     click.echo(
         f"Visit the {project_name} directory and follow the next steps in the Getting started guide (https://docs.flyte.org/en/latest/getting_started.html) to proceed."

setup.py

@@ -68,7 +68,6 @@
         "docstring-parser>=0.9.0",
         "diskcache>=5.2.1",
         "cloudpickle>=2.0.0",
-        "cookiecutter>=1.7.3",
         "numpy",
         "gitpython",
         "kubernetes>=12.0.1",