yt-dlp-stream-to-s3/yt_dlp_stream_to_s3/main.py

import subprocess
import yt_dlp
import boto3
from urllib.error import HTTPError
from botocore.exceptions import ClientError
from yt_dlp_stream_to_s3.client import s3, create_presigned_url
from yt_dlp_stream_to_s3.utils import remove_queries_from_url
from yt_dlp_stream_to_s3.config import AWS_S3_BUCKET_NAME
from yt_dlp_stream_to_s3.errors import YtDlpStreamToS3Error


def yt_dlp_stream_to_s3(
    url: str,
    expiration=86400 * 7,  # 7 days
    s3_extra_args: dict = (),
) -> None:
    media_url = remove_queries_from_url(url)

    ydl_opts = {
        "outtmpl": "%(title)s-%(id)s.%(ext)s",
    }

    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info_dict = ydl.extract_info(media_url, download=False)

            video_title = info_dict.get("title", None)
            video_id = info_dict.get("id", None)
            video_ext = info_dict.get("ext", None)

            filename = f"{video_title}-{video_id}.{video_ext}"

            yt_dlp_process = subprocess.Popen(
                [
                    "yt-dlp",
                    media_url,
                    "--downloader",
                    "ffmpeg",
                    "--hls-use-mpegts",
                    "--no-part",
                    "-o",
                    "-",
                ],
                stdout=subprocess.PIPE,
            )

            s3.upload_fileobj(
                yt_dlp_process.stdout,
                AWS_S3_BUCKET_NAME,
                filename,
                ExtraArgs=s3_extra_args,
            )

            yt_dlp_process.wait()

            return create_presigned_url(AWS_S3_BUCKET_NAME, filename, expiration)
    except HTTPError:
        raise YtDlpStreamToS3Error("Could not connect to the media resource.")

    except yt_dlp.utils.ExtractorError:
        raise YtDlpStreamToS3Error("Could not extra info from media.")

    except yt_dlp.utils.DownloadError:
        raise YtDlpStreamToS3Error("Could not download media. Check logs.")

    except ClientError:
        raise YtDlpStreamToS3Error("Error communicating with the S3 service.")