s3_tree.py

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import List, Dict

import boto3


@dataclass
class S3TreeItem:
    title: str
    path: str
    is_dir: bool
    children: List[S3TreeItem]


class RealS3Service:
    def __init__(self, bucket: str):
        self.s3 = boto3.resource('s3')
        self.bucket = self.s3.Bucket(name=bucket)
    def _put_item(self, root, key):
        *path_items, name = key.split("/")
        dirkey = ""
        dir = root
        for dirname in path_items:
            dirkey += dirname + "/"
            dir = dir.setdefault(dirkey, {})
        dir[key] = None

    def _walk_tree(self, bucket_name, tree: Dict):
        for key, value in tree.items():
            *_, name, end = key.rsplit("/")
            if isinstance(value, dict):
                children = list(self._walk_tree(bucket_name, value))
                dir = S3TreeItem(name, f"s3://{bucket_name}/{key}", True, children)
                yield dir
            else:
                yield S3TreeItem(name, f"s3://{bucket_name}/{key}", False, [])

    def walk(self):
        items = self.bucket.objects.all()
        root = {}
        for item in items:
            self._put_item(root, item.key)

        return S3TreeItem("/", f"s3://{self.bucket.name}/", True, list(self._walk_tree(self.bucket.name, root)))