nebu 0.1.24__py3-none-any.whl → 0.1.29__py3-none-any.whl

nebu/data.py

@@ -0,0 +1,855 @@
+import os
+import subprocess
+from datetime import datetime, timedelta, timezone
+from typing import Any, Dict, List, Optional, Tuple
+from urllib.parse import urlparse
+
+import boto3
+from botocore.exceptions import ClientError
+
+
+def rclone_copy(
+    source_dir: str,
+    destination: str,
+    dry_run: bool = False,
+    transfers: int = 4,
+    extra_args: Optional[List[str]] = None,
+    verbose: bool = True,
+) -> bool:
+    """
+    Upload a directory to a remote bucket using `rclone copy`.
+
+    Args:
+        source_dir (str): Path to local directory to upload.
+        destination (str): Remote destination, e.g., 's3:my-bucket/path'.
+        dry_run (bool): If True, performs a dry run without uploading.
+        transfers (int): Number of parallel transfers.
+        extra_args (Optional[List[str]]): Additional rclone flags.
+        verbose (bool): If True, prints command and output live.
+
+    Returns:
+        bool: True if upload succeeded, False otherwise.
+    """
+    command = [
+        "rclone",
+        "copy",
+        source_dir,
+        destination,
+        f"--transfers={transfers}",
+        "--progress",
+    ]
+
+    if dry_run:
+        command.append("--dry-run")
+    if extra_args:
+        command.extend(extra_args)
+
+    if verbose:
+        print("Running command:", " ".join(command))
+
+    try:
+        process = subprocess.Popen(
+            command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
+        )
+
+        if not process.stdout:
+            raise Exception("No output from rclone")
+
+        for line in process.stdout:
+            if verbose:
+                print(line.strip())
+
+        return process.wait() == 0
+
+    except Exception as e:
+        print(f"Error during rclone copy: {e}")
+        return False
+
+
+def find_latest_checkpoint(training_dir: str) -> Optional[str]:
+    """
+    Finds the checkpoint directory with the highest step number in a Hugging Face
+    training output directory.
+
+    Args:
+        training_dir (str): The path to the training output directory.
+
+    Returns:
+        Optional[str]: The path to the latest checkpoint directory, or None if
+                       no checkpoint directories are found or the directory
+                       doesn't exist.
+    """
+    latest_step = -1
+    latest_checkpoint_dir = None
+
+    if not os.path.isdir(training_dir):
+        print(f"Error: Directory not found: {training_dir}")
+        return None
+
+    for item in os.listdir(training_dir):
+        item_path = os.path.join(training_dir, item)
+        if os.path.isdir(item_path) and item.startswith("checkpoint-"):
+            try:
+                step_str = item.split("-")[-1]
+                if step_str.isdigit():
+                    step = int(step_str)
+                    if step > latest_step:
+                        latest_step = step
+                        latest_checkpoint_dir = item_path
+            except (ValueError, IndexError):
+                # Ignore items that don't match the expected pattern
+                continue
+
+    return latest_checkpoint_dir
+
+
+def _parse_s3_path(path: str) -> Tuple[Optional[str], Optional[str]]:
+    """Standalone helper: Parses an S3 path (s3://bucket/prefix) into bucket and prefix."""
+    parsed = urlparse(path)
+    if parsed.scheme != "s3":
+        return None, None
+    bucket = parsed.netloc
+    prefix = parsed.path.lstrip("/")
+    return bucket, prefix
+
+
+class Bucket:
+    """Handles interactions with AWS S3."""
+
+    def __init__(self, verbose: bool = True):
+        """
+        Initializes the S3 handler.
+
+        Args:
+            verbose (bool): If True, prints status messages. Defaults to True.
+        """
+        self.client = boto3.client("s3")
+        self.verbose = verbose
+
+    def _parse_path(self, path: str) -> Tuple[Optional[str], Optional[str]]:
+        """Class method: Parses an S3 path (s3://bucket/prefix) into bucket and prefix."""
+        # Reusing the standalone logic here for consistency
+        return _parse_s3_path(path)
+
+    def _list_objects(
+        self, bucket: str, prefix: Optional[str]
+    ) -> Dict[str, Dict[str, Any]]:
+        """Class method: Lists objects in an S3 prefix."""
+        objects: Dict[str, Dict[str, Any]] = {}
+        paginator = self.client.get_paginator("list_objects_v2")
+        list_prefix = prefix or ""
+        if self.verbose:
+            print(f"Listing objects in s3://{bucket}/{list_prefix}...")
+
+        operation_parameters = {"Bucket": bucket}
+        if list_prefix:
+            operation_parameters["Prefix"] = list_prefix
+
+        try:
+            page_iterator = paginator.paginate(**operation_parameters)
+            for page in page_iterator:
+                if "Contents" in page:
+                    for obj in page["Contents"]:
+                        if obj["Key"].endswith("/") and obj["Size"] == 0:
+                            continue
+                        relative_key: Optional[str] = None
+                        current_prefix = prefix or ""
+                        if current_prefix and obj["Key"].startswith(current_prefix):
+                            prefix_adjusted = current_prefix + (
+                                "" if current_prefix.endswith("/") else "/"
+                            )
+                            if obj["Key"] == current_prefix.rstrip("/"):
+                                relative_key = os.path.basename(obj["Key"])
+                            elif obj["Key"].startswith(prefix_adjusted):
+                                relative_key = obj["Key"][len(prefix_adjusted) :]
+                            else:
+                                potential_rel_key = obj["Key"][len(current_prefix) :]
+                                relative_key = potential_rel_key.lstrip("/")
+                        elif not current_prefix:
+                            relative_key = obj["Key"]
+                        if not relative_key:
+                            continue
+                        last_modified = obj["LastModified"]
+                        if last_modified.tzinfo is None:
+                            last_modified = last_modified.replace(tzinfo=timezone.utc)
+                        objects[relative_key] = {
+                            "path": f"s3://{bucket}/{obj['Key']}",
+                            "key": obj["Key"],
+                            "size": obj["Size"],
+                            "mtime": last_modified,
+                            "type": "s3",
+                        }
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchBucket":
+                if self.verbose:
+                    print(f"Error: Bucket '{bucket}' not found.")
+            elif e.response["Error"]["Code"] == "NoSuchKey" and prefix:
+                if self.verbose:
+                    print(
+                        f"Prefix s3://{bucket}/{prefix} not found (treating as empty)."
+                    )
+            else:
+                print(f"Error listing S3 objects: {e}")
+            if e.response["Error"]["Code"] == "NoSuchBucket":
+                return {}
+        except Exception as e:
+            print(f"An unexpected error occurred listing S3 objects: {e}")
+            return {}
+        if self.verbose:
+            print(f"Found {len(objects)} objects in S3.")
+        return objects
+
+    def _list_local(self, local_dir: str) -> Dict[str, Dict[str, Any]]:
+        """Class method: Lists files in a local directory."""
+        files: Dict[str, Dict[str, Any]] = {}
+        if not os.path.exists(local_dir):
+            if self.verbose:
+                print(
+                    f"Warning: Local path not found: {local_dir} (treating as empty)."
+                )
+            return files
+        if os.path.isfile(local_dir):
+            if self.verbose:
+                print(
+                    f"Warning: Source {local_dir} is a file, not a directory. Syncing single file."
+                )
+            try:
+                file_name = os.path.basename(local_dir)
+                files[file_name] = {
+                    "path": local_dir,
+                    "size": os.path.getsize(local_dir),
+                    "mtime": datetime.fromtimestamp(
+                        os.path.getmtime(local_dir), tz=timezone.utc
+                    ),
+                    "type": "local",
+                }
+            except OSError as e:
+                print(f"Error accessing source file {local_dir}: {e}")
+            return files
+        if self.verbose:
+            print(f"Scanning local directory: {local_dir}...")
+        for root, _, file_list in os.walk(local_dir):
+            for file_name in file_list:
+                local_path = os.path.join(root, file_name)
+                try:
+                    relative_path = os.path.relpath(local_path, local_dir).replace(
+                        "\\", "/"
+                    )
+                    files[relative_path] = {
+                        "path": local_path,
+                        "size": os.path.getsize(local_path),
+                        "mtime": datetime.fromtimestamp(
+                            os.path.getmtime(local_path), tz=timezone.utc
+                        ),
+                        "type": "local",
+                    }
+                except OSError as e:
+                    print(f"Warning: Could not get metadata for {local_path}: {e}")
+                except Exception as e:
+                    print(f"Warning: Unexpected error processing {local_path}: {e}")
+        if self.verbose:
+            print(f"Found {len(files)} files locally.")
+        return files
+
+    def sync(
+        self,
+        source: str,
+        destination: str,
+        delete: bool = False,
+        dry_run: bool = False,
+    ) -> None:
+        """
+        Synchronizes files between a source and a destination (local or S3).
+        Compares file sizes and modification times. Copies if missing, larger, or newer.
+        Optionally deletes extraneous files from the destination.
+        Args:
+            source (str): The source path (local directory/file or s3://...).
+            destination (str): The destination path (local directory or s3://...).
+            delete (bool): If True, delete extraneous files from the destination.
+            dry_run (bool): If True, print actions without performing them.
+        """
+        mtime_tolerance = timedelta(seconds=2)
+        src_bucket, src_prefix = self._parse_path(source)
+        dest_bucket, dest_prefix = self._parse_path(destination)
+        source_items: Dict[str, Dict[str, Any]] = {}
+        dest_items: Dict[str, Dict[str, Any]] = {}
+        sync_direction = ""
+        is_single_file_sync = False
+
+        if src_bucket is None and dest_bucket is not None:
+            sync_direction = "upload"
+            source_items = self._list_local(source)
+            dest_items = self._list_objects(dest_bucket, dest_prefix)
+            if not source_items and not os.path.exists(source):
+                print(
+                    f"Error: Source path {source} not found and is not empty."
+                )  # Check needed? list_local handles it.
+                # return # Let it proceed if source is just empty
+            if os.path.isfile(source):
+                is_single_file_sync = True
+            # current_dest_prefix = dest_prefix or "" # Moved closer to usage
+
+        elif src_bucket is not None and dest_bucket is None:
+            sync_direction = "download"
+            source_items = self._list_objects(src_bucket, src_prefix)
+            if os.path.exists(destination) and not os.path.isdir(destination):
+                print(
+                    f"Error: Local destination '{destination}' exists but is not a directory."
+                )
+                return
+            dest_items = self._list_local(destination)
+            if not dry_run:
+                os.makedirs(destination, exist_ok=True)
+            elif not os.path.isdir(destination) and self.verbose:
+                print(f"Dry run: Would create local directory {destination}")
+
+        elif src_bucket is None and dest_bucket is None:
+            print(
+                "Error: Both source and destination are local paths. Use standard file copy tools."
+            )
+            return
+        elif src_bucket is not None and dest_bucket is not None:
+            print(
+                "Error: S3 to S3 sync not implemented. Use AWS CLI or S3 Batch Operations."
+            )
+            return
+        else:
+            print("Error: Invalid source or destination path combination.")
+            return
+
+        actions_to_perform: List[Dict[str, Any]] = []
+        source_keys = set(source_items.keys())
+        dest_keys = set(dest_items.keys())
+
+        for rel_key in source_keys:
+            src_item = source_items[rel_key]
+            dest_item = dest_items.get(rel_key)
+            reason = ""
+            if dest_item is None:
+                reason = "does not exist in destination"
+            else:
+                if src_item["size"] != dest_item["size"]:
+                    reason = f"size differs (src: {src_item['size']}, dest: {dest_item['size']})"
+                elif src_item["mtime"] > (dest_item["mtime"] + mtime_tolerance):
+                    reason = f"is newer in source (src: {src_item['mtime']}, dest: {dest_item['mtime']})"
+            if reason:
+                action_type = "upload" if sync_direction == "upload" else "download"
+                dest_full_path_or_key: Optional[str] = None
+                if sync_direction == "upload":
+                    # Define current_dest_prefix here, just before use
+                    current_dest_prefix = dest_prefix or ""
+                    final_dest_key = (
+                        rel_key
+                        if is_single_file_sync
+                        else os.path.join(current_dest_prefix, rel_key).replace(
+                            "\\", "/"
+                        )
+                    )
+                    if not current_dest_prefix and final_dest_key.startswith("/"):
+                        final_dest_key = final_dest_key.lstrip("/")
+                    dest_full_path_or_key = f"s3://{dest_bucket}/{final_dest_key}"
+                else:
+                    dest_full_path_or_key = os.path.join(
+                        destination, rel_key.replace("/", os.sep)
+                    )
+                actions_to_perform.append(
+                    {
+                        "action": action_type,
+                        "relative_key": rel_key,
+                        "source_path": src_item["path"],
+                        "source_mtime": src_item.get("mtime"),
+                        "dest_full_path_or_key": dest_full_path_or_key,
+                        "dest_bucket": dest_bucket,
+                        "dest_prefix": dest_prefix,
+                        "s3_key_full_src": src_item.get("key")
+                        if sync_direction == "download"
+                        else None,
+                        "source_bucket": src_bucket,
+                        "reason": reason,
+                    }
+                )
+
+        if delete:
+            keys_to_delete = dest_keys - source_keys
+            for rel_key in keys_to_delete:
+                dest_item = dest_items[rel_key]
+                action_type = (
+                    "delete_s3" if sync_direction == "upload" else "delete_local"
+                )
+                actions_to_perform.append(
+                    {
+                        "action": action_type,
+                        "relative_key": rel_key,
+                        "path_to_delete": dest_item["path"],
+                        "s3_key_full_dest": dest_item.get("key")
+                        if sync_direction == "upload"
+                        else None,
+                        "dest_bucket": dest_bucket,
+                        "reason": "does not exist in source",
+                    }
+                )
+
+        uploads_done = downloads_done = deletions_done = 0
+        s3_deletions_batch: List[Dict[str, str]] = []
+        if not actions_to_perform:
+            if self.verbose:
+                print("Source and destination are already synchronized.")
+            # Optional: Add check if source exists if sync_direction == "upload" and not os.path.exists(source):
+            return
+
+        for action in actions_to_perform:
+            reason = action["reason"]
+            dest_full_path_or_key = action["dest_full_path_or_key"]
+            if action["action"] == "upload":
+                local_path = action["source_path"]
+                if not isinstance(dest_full_path_or_key, str):
+                    print(f"ERROR: Invalid dest path: {dest_full_path_or_key}")
+                    continue
+                _, upload_key = self._parse_path(dest_full_path_or_key)
+                target_bucket = action["dest_bucket"]
+                if self.verbose:
+                    print(f"Upload: {local_path} to {dest_full_path_or_key} ({reason})")
+                if not dry_run:
+                    if target_bucket and upload_key is not None:
+                        try:
+                            self.client.upload_file(
+                                local_path, target_bucket, upload_key
+                            )
+                            uploads_done += 1
+                        except ClientError as e:
+                            print(f"ERROR uploading {local_path}: {e}")
+                        except Exception as e:
+                            print(f"ERROR uploading {local_path}: {e}")
+                    else:
+                        print(
+                            f"ERROR: Invalid S3 target: bucket={target_bucket}, key={upload_key}"
+                        )
+            elif action["action"] == "download":
+                s3_key_full = action["s3_key_full_src"]
+                local_path = dest_full_path_or_key
+                source_bucket_dl = action["source_bucket"]
+                if self.verbose:
+                    print(
+                        f"Download: {action['source_path']} to {local_path} ({reason})"
+                    )
+                if not isinstance(local_path, str):
+                    print(f"ERROR: Invalid local dest path: {local_path}")
+                    continue
+                if not dry_run:
+                    if source_bucket_dl and s3_key_full and local_path:
+                        try:
+                            local_file_dir = os.path.dirname(local_path)
+                            os.makedirs(local_file_dir, exist_ok=True)
+                            self.client.download_file(
+                                source_bucket_dl, s3_key_full, local_path
+                            )
+                            downloads_done += 1
+                        except ClientError as e:
+                            print(f"ERROR downloading {s3_key_full}: {e}")
+                        except OSError as e:
+                            print(f"ERROR creating/writing {local_path}: {e}")
+                        except Exception as e:
+                            print(f"ERROR downloading {s3_key_full}: {e}")
+                    else:
+                        print(
+                            f"ERROR: Invalid download params: bucket={source_bucket_dl}, key={s3_key_full}, local={local_path}"
+                        )
+            elif action["action"] == "delete_s3":
+                s3_key_to_delete = action["s3_key_full_dest"]
+                target_bucket_del = action["dest_bucket"]
+                if target_bucket_del and s3_key_to_delete:
+                    if self.verbose:
+                        print(f"Delete S3: {action['path_to_delete']} ({reason})")
+                    if isinstance(s3_key_to_delete, str):
+                        s3_deletions_batch.append({"Key": s3_key_to_delete})
+                    else:
+                        print(f"ERROR: Invalid S3 key for deletion: {s3_key_to_delete}")
+                else:
+                    print(
+                        f"ERROR: Invalid S3 target for deletion: bucket={target_bucket_del}, key={s3_key_to_delete}"
+                    )
+            elif action["action"] == "delete_local":
+                local_path_to_delete = action["path_to_delete"]
+                if self.verbose:
+                    print(f"Delete Local: {local_path_to_delete} ({reason})")
+                if not dry_run:
+                    try:
+                        os.remove(local_path_to_delete)
+                        deletions_done += 1
+                    except OSError as e:
+                        print(f"ERROR deleting local file {local_path_to_delete}: {e}")
+
+        if s3_deletions_batch:
+            target_bucket_del_batch = next(
+                (
+                    a["dest_bucket"]
+                    for a in actions_to_perform
+                    if a["action"] == "delete_s3"
+                ),
+                None,
+            )
+            if not dry_run and target_bucket_del_batch:
+                deleted_count_batch = 0
+                for i in range(0, len(s3_deletions_batch), 1000):
+                    batch = s3_deletions_batch[i : i + 1000]
+                    delete_payload = {"Objects": batch, "Quiet": False}
+                    try:
+                        response = self.client.delete_objects(
+                            Bucket=target_bucket_del_batch, Delete=delete_payload
+                        )
+                        deleted_count_batch += len(batch)
+                        if "Errors" in response and response["Errors"]:
+                            deleted_count_batch -= len(response["Errors"])
+                            for error in response["Errors"]:
+                                print(
+                                    f"ERROR deleting S3 object {error['Key']}: {error['Code']} - {error['Message']}"
+                                )
+                    except ClientError as e:
+                        print(f"ERROR deleting S3 objects batch: {e}")
+                        deleted_count_batch = 0
+                    except Exception as e:
+                        print(f"ERROR deleting S3 objects batch: {e}")
+                        deleted_count_batch = 0
+                deletions_done += deleted_count_batch
+            elif target_bucket_del_batch:
+                deletions_done = len(s3_deletions_batch)
+            else:
+                print(
+                    "Warning: Could not determine target bucket for S3 deletion batch."
+                )
+
+        if dry_run:
+            if self.verbose:
+                upload_count = sum(
+                    1 for a in actions_to_perform if a["action"] == "upload"
+                )
+                download_count = sum(
+                    1 for a in actions_to_perform if a["action"] == "download"
+                )
+                delete_s3_count = len(s3_deletions_batch)
+                delete_local_count = sum(
+                    1 for a in actions_to_perform if a["action"] == "delete_local"
+                )
+                print("\n--- DRY RUN SUMMARY ---")
+                if sync_direction == "upload":
+                    print(f"Would upload: {upload_count} file(s)")
+                    if delete:
+                        print(f"Would delete from S3: {delete_s3_count} object(s)")
+                elif sync_direction == "download":
+                    print(f"Would download: {download_count} file(s)")
+                    if delete:
+                        print(f"Would delete locally: {delete_local_count} file(s)")
+                print("--- END DRY RUN ---")
+        else:
+            if self.verbose:
+                if sync_direction == "upload":
+                    print(
+                        f"Sync completed. Uploaded: {uploads_done} file(s). Deleted from S3: {deletions_done if delete else 0} object(s)."
+                    )
+                elif sync_direction == "download":
+                    print(
+                        f"Sync completed. Downloaded: {downloads_done} file(s). Deleted locally: {deletions_done if delete else 0} file(s)."
+                    )
+
+    def copy(
+        self,
+        source: str,
+        destination: str,
+    ) -> None:
+        """
+        Copies files or directories between local paths and S3 URIs.
+        Handles:
+        - Local file to S3 object
+        - Local directory to S3 prefix (recursive)
+        - S3 object to local file
+        - S3 prefix to local directory (recursive)
+        Does NOT handle:
+        - Local to Local (use shutil)
+        - S3 to S3 (use AWS CLI or boto3 object copy)
+        Args:
+            source (str): The source path (local file/dir or s3://...).
+            destination (str): The destination path (local file/dir or s3://...).
+        """
+        src_bucket, src_prefix = self._parse_path(source)
+        dest_bucket, dest_prefix = self._parse_path(destination)
+
+        if src_bucket is None and dest_bucket is None:
+            print(
+                "Error: Both source and destination are local. Use 'shutil.copy' or 'shutil.copytree'."
+            )
+            return
+        if src_bucket is not None and dest_bucket is not None:
+            print(
+                "Error: S3 to S3 copy not implemented. Use 'aws s3 cp' or boto3 'copy_object'."
+            )
+            return
+
+        # Upload: Local to S3
+        if src_bucket is None and dest_bucket is not None:
+            if not os.path.exists(source):
+                print(f"Error: Local source path not found: {source}")
+                return
+            current_dest_prefix = dest_prefix or ""
+
+            if os.path.isfile(source):
+                if not current_dest_prefix or destination.endswith("/"):
+                    s3_key = os.path.join(
+                        current_dest_prefix, os.path.basename(source)
+                    ).replace("\\", "/")
+                else:
+                    s3_key = current_dest_prefix
+                if self.verbose:
+                    print(f"Uploading {source} to s3://{dest_bucket}/{s3_key}")
+                try:
+                    self.client.upload_file(source, dest_bucket, s3_key)
+                    if self.verbose:
+                        print("Upload complete.")
+                except ClientError as e:
+                    print(f"ERROR uploading {source}: {e}")
+                except Exception as e:
+                    print(f"ERROR uploading {source}: {e}")
+
+            elif os.path.isdir(source):
+                if self.verbose:
+                    print(
+                        f"Uploading directory {source}/* to s3://{dest_bucket}/{current_dest_prefix}/"
+                    )
+                files_uploaded = files_failed = 0
+                for root, _, files in os.walk(source):
+                    for file in files:
+                        local_path = os.path.join(root, file)
+                        relative_path = os.path.relpath(local_path, source)
+                        s3_key = os.path.join(
+                            current_dest_prefix, relative_path
+                        ).replace("\\", "/")
+                        if self.verbose:
+                            print(
+                                f"  Uploading {local_path} to s3://{dest_bucket}/{s3_key}"
+                            )
+                        try:
+                            self.client.upload_file(local_path, dest_bucket, s3_key)
+                            files_uploaded += 1
+                        except ClientError as e:
+                            print(f"  ERROR uploading {local_path}: {e}")
+                            files_failed += 1
+                        except Exception as e:
+                            print(f"  ERROR uploading {local_path}: {e}")
+                            files_failed += 1
+                if self.verbose:
+                    print(
+                        f"Directory upload complete. Files uploaded: {files_uploaded}, Failed: {files_failed}"
+                    )
+            else:
+                print(f"Error: Source {source} is neither a file nor a directory.")
+
+        # Download: S3 to Local
+        elif src_bucket is not None and dest_bucket is None:
+            is_prefix_download = False
+            single_object_key = None
+            current_src_prefix = src_prefix or ""  # Ensure not None
+
+            if source.endswith("/"):
+                is_prefix_download = True
+            else:
+                try:
+                    if current_src_prefix:
+                        self.client.head_object(
+                            Bucket=src_bucket, Key=current_src_prefix
+                        )
+                        single_object_key = current_src_prefix
+                    else:
+                        # Path like s3://bucket, treat as prefix download
+                        is_prefix_download = True
+                except ClientError as e:
+                    if e.response["Error"]["Code"] == "404":
+                        is_prefix_download = True  # Assume prefix if object not found
+                    elif e.response["Error"]["Code"] == "NoSuchBucket":
+                        print(f"Error: Source bucket '{src_bucket}' not found.")
+                        return
+                    else:
+                        print(
+                            f"Error checking S3 source s3://{src_bucket}/{current_src_prefix}: {e}"
+                        )
+                        return
+                except Exception as e:
+                    print(
+                        f"Error checking S3 source s3://{src_bucket}/{current_src_prefix}: {e}"
+                    )
+                    return
+
+            if single_object_key is not None:
+                if os.path.isdir(destination) or destination.endswith(os.sep):
+                    local_dest_path = os.path.join(
+                        destination, os.path.basename(single_object_key)
+                    )
+                    os.makedirs(destination, exist_ok=True)
+                else:
+                    local_dest_path = destination
+                    parent_dir = os.path.dirname(local_dest_path)
+                    if parent_dir:
+                        os.makedirs(parent_dir, exist_ok=True)
+                if self.verbose:
+                    print(
+                        f"Downloading s3://{src_bucket}/{single_object_key} to {local_dest_path}"
+                    )
+                try:
+                    self.client.download_file(
+                        src_bucket, single_object_key, local_dest_path
+                    )
+                    if self.verbose:
+                        print("Download complete.")
+                except ClientError as e:
+                    print(f"ERROR downloading {single_object_key}: {e}")
+                except OSError as e:
+                    print(f"ERROR creating/writing {local_dest_path}: {e}")
+                except Exception as e:
+                    print(f"ERROR downloading {single_object_key}: {e}")
+
+            elif is_prefix_download:
+                if os.path.exists(destination) and not os.path.isdir(destination):
+                    print(
+                        f"Error: Local destination '{destination}' exists but is not a directory."
+                    )
+                    return
+                os.makedirs(destination, exist_ok=True)
+                if self.verbose:
+                    print(
+                        f"Downloading prefix s3://{src_bucket}/{current_src_prefix}/* to {destination}/"
+                    )
+                paginator = self.client.get_paginator("list_objects_v2")
+                files_downloaded = files_failed = 0
+                operation_parameters = {"Bucket": src_bucket}
+                # The problematic line for the linter, re-adding type ignore
+                if current_src_prefix:  # type: ignore
+                    operation_parameters["Prefix"] = current_src_prefix
+                try:
+                    page_iterator = paginator.paginate(**operation_parameters)
+                    found_objects = False
+                    for page in page_iterator:
+                        if "Contents" in page:
+                            found_objects = True
+                            for obj in page["Contents"]:
+                                s3_key = obj["Key"]
+                                if s3_key.endswith("/") and obj["Size"] == 0:
+                                    continue
+                                relative_key = s3_key
+                                if current_src_prefix:
+                                    if s3_key.startswith(current_src_prefix):
+                                        if s3_key == current_src_prefix.rstrip("/"):
+                                            relative_key = os.path.basename(s3_key)
+                                        else:
+                                            prefix_adjusted = current_src_prefix + (
+                                                ""
+                                                if current_src_prefix.endswith("/")
+                                                else "/"
+                                            )
+                                            if s3_key.startswith(prefix_adjusted):
+                                                relative_key = s3_key[
+                                                    len(prefix_adjusted) :
+                                                ]
+                                            elif not current_src_prefix.endswith("/"):
+                                                relative_key = s3_key[
+                                                    len(current_src_prefix) :
+                                                ].lstrip("/")
+                                if not relative_key:
+                                    continue
+                                local_dest_path = os.path.join(
+                                    destination, relative_key.replace("/", os.sep)
+                                )
+                                local_dest_dir = os.path.dirname(local_dest_path)
+                                if self.verbose:
+                                    print(
+                                        f"  Downloading s3://{src_bucket}/{s3_key} to {local_dest_path}"
+                                    )
+                                try:
+                                    if local_dest_dir:
+                                        os.makedirs(local_dest_dir, exist_ok=True)
+                                    self.client.download_file(
+                                        src_bucket, s3_key, local_dest_path
+                                    )
+                                    files_downloaded += 1
+                                except ClientError as e:
+                                    print(f"  ERROR downloading {s3_key}: {e}")
+                                    files_failed += 1
+                                except OSError as e:
+                                    print(
+                                        f"  ERROR creating/writing {local_dest_path}: {e}"
+                                    )
+                                    files_failed += 1
+                                except Exception as e:
+                                    print(f"  ERROR downloading {s3_key}: {e}")
+                                    files_failed += 1
+                    if not found_objects and self.verbose:
+                        print(
+                            f"Warning: No objects found at source prefix s3://{src_bucket}/{current_src_prefix}"
+                        )
+                    if self.verbose:
+                        print(
+                            f"Prefix download complete. Files downloaded: {files_downloaded}, Failed: {files_failed}"
+                        )
+                except ClientError as e:
+                    if e.response["Error"]["Code"] == "NoSuchBucket":
+                        print(f"Error: Source bucket '{src_bucket}' not found.")
+                    else:
+                        print(
+                            f"Error listing objects in s3://{src_bucket}/{current_src_prefix}: {e}"
+                        )
+                except Exception as e:
+                    print(
+                        f"Error listing objects in s3://{src_bucket}/{current_src_prefix}: {e}"
+                    )
+        else:
+            print("Error: Unknown copy operation type.")
+
+    def check(self, s3_uri: str) -> bool:
+        """
+        Check if an object or prefix exists in an S3 bucket using an S3 URI.
+
+        Args:
+            s3_uri (str): The S3 URI (e.g., 's3://my-bucket/my-key' or 's3://my-bucket/my-prefix/').
+                          Use a trailing '/' to check for a prefix/directory.
+
+        Returns:
+            bool: True if the object or prefix exists, False otherwise.
+        """
+        # Use the class client and parse method
+        bucket_name, s3_key = self._parse_path(s3_uri)
+
+        if bucket_name is None or s3_key is None:
+            # _parse_path returns None, None if scheme is not 's3'
+            print(f"Error: Invalid S3 URI format: {s3_uri}")
+            return False
+
+        is_prefix = s3_key.endswith("/")
+
+        try:
+            if is_prefix:
+                # Check for prefix existence by listing objects
+                # Handle the case where s3_key might be empty if URI is just s3://bucket/
+                list_prefix = s3_key if s3_key else ""
+                response = self.client.list_objects_v2(
+                    Bucket=bucket_name, Prefix=list_prefix, MaxKeys=1
+                )
+                # Check if any objects OR common prefixes (folders) are returned for the prefix
+                return "Contents" in response or "CommonPrefixes" in response
+            else:
+                # Check for object existence
+                self.client.head_object(Bucket=bucket_name, Key=s3_key)
+                return True
+        except ClientError as e:  # Catch boto3 ClientError first
+            # If head_object returns 404 (NoSuchKey), the object doesn't exist
+            # list_objects_v2 does not raise NoSuchKey for prefixes
+            if e.response["Error"]["Code"] == "404":
+                return False
+            elif e.response["Error"]["Code"] == "NoSuchBucket":
+                if self.verbose:
+                    print(
+                        f"Error: Bucket '{bucket_name}' not found (from URI: {s3_uri})."
+                    )
+                return False
+            # Handle other potential errors like AccessDenied differently if needed
+            print(f"Error checking {s3_uri}: {e}")
+            return False
+        except Exception as e:
+            print(f"An unexpected error occurred checking {s3_uri}: {e}")
+            return False