From 3fcc075b2f5dc69f69023a2b1efa0c9c750fe68a Mon Sep 17 00:00:00 2001
From: Ilias Katsakioris <elikatsis@arrikto.com>
Date: Fri, 17 May 2019 16:20:00 +0300
Subject: [PATCH] SDK/Compiler: Add add_pvolumes() method to ContainerOp

Signed-off-by: Ilias Katsakioris <elikatsis@arrikto.com>
---
 sdk/python/kfp/dsl/_container_op.py | 45 +++++++++++++++++++----------
 1 file changed, 29 insertions(+), 16 deletions(-)

diff --git a/sdk/python/kfp/dsl/_container_op.py b/sdk/python/kfp/dsl/_container_op.py
index cea24d339f9..0cf364d8ab6 100644
--- a/sdk/python/kfp/dsl/_container_op.py
+++ b/sdk/python/kfp/dsl/_container_op.py
@@ -973,22 +973,7 @@ def _decorated(*args, **kwargs):
             self.output = list(self.outputs.values())[0]
 
         self.pvolumes = {}
-        if pvolumes:
-            for mount_path, pvolume in pvolumes.items():
-                if hasattr(pvolume, "dependent_names"): #TODO: Replace with type check
-                    self.dependent_names.extend(pvolume.dependent_names)
-                else:
-                    pvolume = PipelineVolume(volume=pvolume)
-                self.pvolumes[mount_path] = pvolume.after(self)
-                self.add_volume(pvolume)
-                self._container.add_volume_mount(V1VolumeMount(
-                    name=pvolume.name,
-                    mount_path=mount_path
-                ))
-
-        self.pvolume = None
-        if self.pvolumes and len(self.pvolumes) == 1:
-            self.pvolume = list(self.pvolumes.values())[0]
+        self.add_pvolumes(pvolumes)
 
     @property
     def command(self):
@@ -1050,6 +1035,34 @@ def _set_metadata(self, metadata):
             if len(self.outputs) == 1:
                 self.output = list(self.outputs.values())[0]
 
+    def add_pvolumes(self,
+                     pvolumes: Dict[str, V1Volume] = None):
+        """Updates the existing pvolumes dict, extends volumes and volume_mounts
+        and redefines the pvolume attribute.
+
+        Args:
+            pvolumes: Dictionary. Keys are mount paths, values are Kubernetes
+                      volumes or inherited types (e.g. PipelineVolumes).
+        """
+        if pvolumes:
+            for mount_path, pvolume in pvolumes.items():
+                if hasattr(pvolume, "dependent_names"):
+                    self.dependent_names.extend(pvolume.dependent_names)
+                else:
+                    pvolume = PipelineVolume(volume=pvolume)
+                self.pvolumes[mount_path] = pvolume.after(self)
+                self.add_volume(pvolume)
+                self._container.add_volume_mount(V1VolumeMount(
+                    name=pvolume.name,
+                    mount_path=mount_path
+                ))
+
+        self.pvolume = None
+        if len(self.pvolumes) == 1:
+            self.pvolume = list(self.pvolumes.values())[0]
+        return self
+
+
 # proxy old ContainerOp properties to ContainerOp.container
 # with PendingDeprecationWarning.
 ContainerOp = _proxy_container_op_props(ContainerOp)