From 468f44a4c02c79549c47819be01be4d99770f109 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Sun, 11 Nov 2018 20:06:05 -0800 Subject: [PATCH] Simplified GCSHelper by extracting duplicate code --- sdk/python/kfp/compiler/_component_builder.py | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/sdk/python/kfp/compiler/_component_builder.py b/sdk/python/kfp/compiler/_component_builder.py index 7c09e34e9657..b5f2522e8e5a 100644 --- a/sdk/python/kfp/compiler/_component_builder.py +++ b/sdk/python/kfp/compiler/_component_builder.py @@ -29,36 +29,31 @@ class GCSHelper(object): """ GCSHelper manages the connection with the GCS storage """ @staticmethod - def upload_gcs_file(local_path, gcs_path): + def get_blob_from_gcs_uri(gcs_path): + from google.cloud import storage pure_path = PurePath(gcs_path) gcs_bucket = pure_path.parts[1] gcs_blob = '/'.join(pure_path.parts[2:]) client = storage.Client() - bucket = client.get_bucket(gcs_bucket) + bucket = client.bucket(gcs_bucket) blob = bucket.blob(gcs_blob) + return blob + + @staticmethod + def upload_gcs_file(local_path, gcs_path): + blob = GCSHelper.get_blob_from_gcs_uri(gcs_path) blob.upload_from_filename(local_path) @staticmethod def remove_gcs_blob(gcs_path): - pure_path = PurePath(gcs_path) - gcs_bucket = pure_path.parts[1] - gcs_blob = '/'.join(pure_path.parts[2:]) - client = storage.Client() - bucket = client.get_bucket(gcs_bucket) - blob = bucket.blob(gcs_blob) + blob = GCSHelper.get_blob_from_gcs_uri(gcs_path) blob.delete() @staticmethod def download_gcs_blob(local_path, gcs_path): - pure_path = PurePath(gcs_path) - gcs_bucket = pure_path.parts[1] - gcs_blob = '/'.join(pure_path.parts[2:]) - client = storage.Client() - bucket = client.get_bucket(gcs_bucket) - blob = bucket.blob(gcs_blob) + blob = GCSHelper.get_blob_from_gcs_uri(gcs_path) blob.download_to_filename(local_path) - class DockerfileHelper(object): """ Dockerfile Helper generates a tarball with dockerfile, ready for docker build arc_dockerfile_name: dockerfile filename that is stored in the tarball