feat: add optional followlinks to datasamples spec (#426)

Signed-off-by: ThibaultFy <[email protected]>

Author: ThibaultFy

substra/sdk/schemas.py

@@ -1,6 +1,7 @@
 import contextlib
 import enum
 import json
+import logging
 import pathlib
 import typing
 import uuid
@@ -14,6 +15,8 @@
 
 from substra.sdk import utils
 
+logger = logging.getLogger(__name__)
+
 _SERVER_NAMES = {
     "dataset": "data_manager",
     "summary_task": "task",
@@ -142,11 +145,14 @@ class DataSampleSpec(_Spec):
     """Specification to create one or many data samples
     To create one data sample, use the 'path' field, otherwise use
     the 'paths' field.
+    Use 'followlinks' to follow symbolic links recursively. Note that it will lead to infinite
+    loops if a symbolic link points to a parent directory.
     """
 
     path: Optional[pathlib.Path] = None  # Path to the data sample if only one
     paths: Optional[List[pathlib.Path]] = None  # Path to the data samples if several
     data_manager_keys: typing.List[str]
+    followlinks: Optional[bool] = False
 
     type_: typing.ClassVar[Type] = Type.DataSample
 
@@ -186,8 +192,13 @@ def build_request_kwargs(self, local):
         # redefine kwargs builder to handle the local paths
         # Serialize and deserialize to prevent errors eg with pathlib.Path
         data = json.loads(self.model_dump_json(exclude_unset=True))
+        if self.followlinks:
+            logger.warning(
+                "The 'followlinks' option is enabled for your datasample registration. It may lead to infinite loops "
+                "if a symbolic link points to a parent directory."
+            )
         if local:
-            with utils.extract_data_sample_files(data) as (data, files):
+            with utils.extract_data_sample_files(data, followlinks=self.followlinks) as (data, files):
                 yield (data, files)
         else:
             yield data, None

substra/sdk/utils.py

@@ -44,25 +44,25 @@ def extract_files(data, file_attributes):
             f.close()
 
 
-def zip_folder(fp, path):
+def zip_folder(fp, path, followlinks=False):
     zipf = zipfile.ZipFile(fp, "w", zipfile.ZIP_DEFLATED)
-    for root, _, files in os.walk(path):
+    for root, _, files in os.walk(path, followlinks=followlinks):
         for f in files:
             abspath = os.path.join(root, f)
             archive_path = os.path.relpath(abspath, start=path)
             zipf.write(abspath, arcname=archive_path)
     zipf.close()
 
 
-def zip_folder_in_memory(path):
+def zip_folder_in_memory(path, followlinks=False):
     fp = io.BytesIO()
-    zip_folder(fp, path)
+    zip_folder(fp, path, followlinks=followlinks)
     fp.seek(0)
     return fp
 
 
 @contextlib.contextmanager
-def extract_data_sample_files(data):
+def extract_data_sample_files(data, followlinks=False):
     # handle data sample specific case; paths and path cases
     data = copy.deepcopy(data)
 
@@ -81,7 +81,7 @@ def extract_data_sample_files(data):
     for k, f in folders.items():
         if not os.path.isdir(f):
             raise exceptions.LoadDataException(f"Paths '{f}' is not an existing directory")
-        files[k] = zip_folder_in_memory(f)
+        files[k] = zip_folder_in_memory(f, followlinks=followlinks)
 
     try:
         yield (data, files)