add replace_filenames argument to wr.s3.copy_objects() #215

Author: igorborgest

awswrangler/s3.py

@@ -2284,6 +2284,7 @@ def copy_objects(
     paths: List[str],
     source_path: str,
     target_path: str,
+    replace_filenames: Optional[Dict[str, str]] = None,
     use_threads: bool = True,
     boto3_session: Optional[boto3.Session] = None,
 ) -> List[str]:
@@ -2334,6 +2335,15 @@ def copy_objects(
     for path in paths:
         path_wo_prefix: str = path.replace(f"{source_path}/", "")
         path_final: str = f"{target_path}/{path_wo_prefix}"
+        if replace_filenames is not None:
+            parts: List[str] = path_final.rsplit(sep="/", maxsplit=1)
+            if len(parts) == 2:
+                path_wo_filename: str = parts[0]
+                filename: str = parts[1]
+                if filename in replace_filenames:
+                    new_filename: str = replace_filenames[filename]
+                    _logger.debug("Replacing filename: %s -> %s", filename, new_filename)
+                    path_final = f"{path_wo_filename}/{new_filename}"
         new_objects.append(path_final)
         batch.append((path, path_final))
     _logger.debug("len(new_objects): %s", len(new_objects))

testing/test_awswrangler/test_data_lake.py

@@ -1348,3 +1348,22 @@ def test_catalog_versioning(bucket, database):
     # Cleaning Up
     wr.catalog.delete_table_if_exists(database=database, table=table)
     wr.s3.delete_objects(path=path)
+
+
+def test_copy_replacing_filename(bucket):
+    path = f"s3://{bucket}/test_copy_replacing_filename/"
+    wr.s3.delete_objects(path=path)
+    df = pd.DataFrame({"c0": [1, 2]})
+    file_path = f"{path}myfile.parquet"
+    wr.s3.to_parquet(df=df, path=file_path)
+    wr.s3.wait_objects_exist(paths=[file_path], use_threads=False)
+    path2 = f"s3://{bucket}/test_copy_replacing_filename2/"
+    wr.s3.copy_objects(
+        paths=[file_path], source_path=path, target_path=path2, replace_filenames={"myfile.parquet": "myfile2.parquet"}
+    )
+    expected_file = f"{path2}myfile2.parquet"
+    wr.s3.wait_objects_exist(paths=[expected_file], use_threads=False)
+    objs = wr.s3.list_objects(path=path2)
+    assert objs[0] == expected_file
+    wr.s3.delete_objects(path=path)
+    wr.s3.delete_objects(path=path2)