Support different modes

Support a flat structure of the downloaded objects

sucket/__init__.py

@@ -6,7 +6,13 @@
 
 
 async def main(
-    loop, bucket_name: str, prefix: str, quiet: bool, skip_prompt: bool, semaphores: int
+    loop,
+    bucket_name: str,
+    prefix: str,
+    quiet: bool,
+    skip_prompt: bool,
+    semaphores: int,
+    mode: str,
 ):
     bucket = Bucket(
         bucket_name=bucket_name,
@@ -15,14 +21,21 @@ async def main(
         quiet=quiet,
         skip_prompt=skip_prompt,
     )
-    await bucket.download_all_objects(prefix)
+    await bucket.download_all_objects(prefix, mode)
 
 
 @click.command()
 @click.argument("bucket_name", type=str)
 @click.argument(
     "prefix", required=False, default="", type=str,
 )
+@click.option(
+    "-m",
+    "--mode",
+    help="The structure to download the objects in.",
+    type=click.Choice(["folder", "flat", "keys-only"], case_sensitive=False),
+    default="folder",
+)
 @click.option(
     "-y", "--yes", is_flag=True, help="Don't prompt for continuing", type=bool
 )
@@ -40,11 +53,22 @@ async def main(
     type=int,
     default=1000,
 )
-def sucket(bucket_name: str, prefix: str, yes: bool, quiet: bool, semaphores: int):
+def sucket(
+    bucket_name: str, prefix: str, yes: bool, quiet: bool, semaphores: int, mode: str
+):
     """ Download all files from a S3 bucket
 
-    By default, everything from a the bucket BUCKET_NAME is downloaded, with an
-    optional key filter, specified with PREFIX
+    Everything from the bucket BUCKET_NAME is downloaded, with an optional key
+    filter, specified with PREFIX
+
+    By default the "folder" mode is used, which will keep the bucket "folder
+    structure" when downloading.
+
+    The "flat" mode will download all objects into the current folder by adding
+    the folder structure into the key name.
+
+    The "keys-only" will completely disregard the folders and put all files in
+    the current folder.
     """
     loop = asyncio.get_event_loop()
     loop.run_until_complete(
@@ -55,5 +79,6 @@ def sucket(bucket_name: str, prefix: str, yes: bool, quiet: bool, semaphores: in
             quiet=quiet,
             skip_prompt=yes or quiet,
             semaphores=semaphores,
+            mode=mode,
         )
     )

sucket/bucket.py

@@ -45,15 +45,24 @@ async def all_objects_paginator(self, prefix: str) -> AsyncIterator[List[Dict]]:
 
                 kwargs["ContinuationToken"] = response["NextContinuationToken"]
 
-    async def _download_object(self, client, obj: Dict) -> int:
+    async def _download_object(self, client, obj: Dict, mode: str) -> int:
         async with self.semaphore:
-            os.makedirs(os.path.dirname(obj["Key"]), exist_ok=True)
+            if mode == "folder":
+                os.makedirs(os.path.dirname(obj["Key"]), exist_ok=True)
+
             if obj["Key"].endswith("/"):
                 # Directory has been created, nothing to download
                 return obj["Size"]
 
+            if mode == "flat":
+                local_file = obj["Key"].replace("/", "-")
+            elif mode == "keys-only":
+                local_file = obj["Key"].rsplit("/", maxsplit=1)[-1]
+            else:
+                local_file = obj["Key"]
+
             res = await client.get_object(Bucket=self.name, Key=obj["Key"])
-            with open(obj["Key"], "wb") as afp:
+            with open(local_file, "wb") as afp:
                 afp.write(await res["Body"].read())
             return obj["Size"]
 
@@ -63,7 +72,7 @@ def secho(self, msg: str, fg: str):
             return
         click.secho(msg, fg=fg)
 
-    async def download_all_objects(self, prefix: str):
+    async def download_all_objects(self, prefix: str, mode: str):
         self.secho("[*] Fetching object metadata...", fg="green")
         objects = []
         try:
@@ -92,7 +101,7 @@ async def download_all_objects(self, prefix: str):
         async with self.session.create_client("s3") as client:
             tasks = []
             for obj in objects:
-                task = asyncio.ensure_future(self._download_object(client, obj))
+                task = asyncio.ensure_future(self._download_object(client, obj, mode))
                 tasks.append(task)
 
             if self.quiet: