fluesterpost/utils.py

from collections import defaultdict
import os

from typing import DefaultDict, Dict, List

from pydub import AudioSegment
import io


def tree() -> DefaultDict:
    return defaultdict(tree)


def get_last_segment(path: str) -> str:
    if path[-1] == '/':
        path = path[:-1]
    return path.split(os.sep)[-1]


FILES_KEY = '_i_files'


def build_file_tree(root_dir: str) -> DefaultDict:
    file_tree = tree()
    root_dir = os.path.normpath(root_dir)  # Normalize the path
    
    for dirpath, dirnames, files in os.walk(root_dir):
        # Get the subdirectory path relative to the root directory
        subdir = os.path.relpath(dirpath, root_dir)
        
        # Split the path into components to navigate the nested dictionary
        path_components = subdir.split(os.sep)

        # Navigate to the current subdirectory in the file tree
        current_subdir = file_tree
        for component in path_components:
            current_subdir = current_subdir[component]

        # Add files to the current subdirectory in the file tree
        current_subdir[FILES_KEY] = files

    return file_tree


# Function to convert defaultdict to dict (for readability)
def defaultdict_to_dict(d: defaultdict) -> dict:
    if isinstance(d, defaultdict):
        d = {k: defaultdict_to_dict(v) for k, v in d.items()}
    return d


def convert_to_mp3(audio_data: bytes, sample_width: int, frame_rate: int, channels: int) -> bytes:
    audio = AudioSegment.from_raw(
        io.BytesIO(audio_data),
        sample_width=sample_width,
        frame_rate=frame_rate,
        channels=channels
    )
    mp3_buffer = io.BytesIO()
    audio.export(mp3_buffer, format="mp3")
    return mp3_buffer.getvalue()