diff --git a/flytekit/clis/sdk_in_container/helpers.py b/flytekit/clis/sdk_in_container/helpers.py index 4c66a2046c..62a8d02edb 100644 --- a/flytekit/clis/sdk_in_container/helpers.py +++ b/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") diff --git a/flytekit/clis/sdk_in_container/init.py b/flytekit/clis/sdk_in_container/init.py index 627b393578..a52e1741ce 100644 --- a/flytekit/clis/sdk_in_container/init.py +++ b/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." diff --git a/setup.py b/setup.py index 5273590d9f..118ab14335 100644 --- a/setup.py +++ b/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",