add filesize to output

main
Juan Olvera 1 year ago
parent e2dd3f4424
commit fed0c2e4fe

17
poetry.lock generated

@ -68,6 +68,17 @@ python-versions = "*"
[package.dependencies]
pycparser = "*"
[[package]]
name = "humanize"
version = "4.6.0"
description = "Python humanize utilities"
category = "main"
optional = false
python-versions = ">=3.7"
[package.extras]
tests = ["freezegun", "pytest", "pytest-cov"]
[[package]]
name = "jmespath"
version = "1.0.1"
@ -173,7 +184,7 @@ websockets = "*"
[metadata]
lock-version = "1.1"
python-versions = "^3.8"
content-hash = "e67744aa22c905b6adc90ef0276bf68b6c7c5e18f32e0ee767bb51ab0146a4a6"
content-hash = "c35992f3e7a627b2df11c502caae24a5af83bda5d02aca973a27cf8a37090201"
[metadata.files]
boto3 = [
@ -370,6 +381,10 @@ cffi = [
{file = "cffi-1.15.1-cp39-cp39-win_amd64.whl", hash = "sha256:70df4e3b545a17496c9b3f41f5115e69a4f2e77e94e1d2a8e1070bc0c38c8a3c"},
{file = "cffi-1.15.1.tar.gz", hash = "sha256:d400bfb9a37b1351253cb402671cea7e89bdecc294e8016a707f6d1d8ac934f9"},
]
humanize = [
{file = "humanize-4.6.0-py3-none-any.whl", hash = "sha256:401201aca462749773f02920139f302450cb548b70489b9b4b92be39fe3c3c50"},
{file = "humanize-4.6.0.tar.gz", hash = "sha256:5f1f22bc65911eb1a6ffe7659bd6598e33dcfeeb904eb16ee1e705a09bf75916"},
]
jmespath = [
{file = "jmespath-1.0.1-py3-none-any.whl", hash = "sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980"},
{file = "jmespath-1.0.1.tar.gz", hash = "sha256:90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe"},

@ -10,6 +10,7 @@ packages = [{include = "yt_dlp_stream_to_s3"}]
python = "^3.8"
yt-dlp = "^2023.3.4"
boto3 = "^1.26.137"
humanize = "^4.6.0"
[build-system]

@ -7,14 +7,15 @@ from yt_dlp_stream_to_s3.config import (
)
from yt_dlp_stream_to_s3.errors import YtDlpStreamToS3ConnectionError
creds = dict(
endpoint_url=AWS_ENDPOINT_URL,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_access_key_id=AWS_ACCESS_KEY_ID,
)
try:
s3 = boto3.client(
"s3",
endpoint_url=AWS_ENDPOINT_URL,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_access_key_id=AWS_ACCESS_KEY_ID,
)
s3_client = boto3.client("s3", **creds)
except Exception:
raise YtDlpStreamToS3ConnectionError(
@ -22,8 +23,13 @@ except Exception:
)
def get_file_size(bucket_name: str, key: str) -> str:
s3_resource = boto3.resource("s3", **creds)
return s3_resource.Object(bucket_name, key).content_length / 1_000_000
def create_presigned_url(bucket_name, object_name, expiration=3600):
return s3.generate_presigned_url(
return s3_client.generate_presigned_url(
"get_object",
Params={"Bucket": bucket_name, "Key": object_name},
ExpiresIn=expiration,

@ -1,9 +1,9 @@
import subprocess
import yt_dlp
import boto3
from typing import Dict
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.client import s3_client, create_presigned_url, get_file_size
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
@ -13,7 +13,7 @@ def yt_dlp_stream_to_s3(
url: str,
expiration=86400 * 7, # 7 days
s3_extra_args: dict = (),
) -> None:
) -> Dict[str, str]:
media_url = remove_queries_from_url(url)
ydl_opts = {
@ -44,7 +44,7 @@ def yt_dlp_stream_to_s3(
stdout=subprocess.PIPE,
)
s3.upload_fileobj(
s3_client.upload_fileobj(
yt_dlp_process.stdout,
AWS_S3_BUCKET_NAME,
filename,
@ -53,7 +53,12 @@ def yt_dlp_stream_to_s3(
yt_dlp_process.wait()
return create_presigned_url(AWS_S3_BUCKET_NAME, filename, expiration)
return dict(
presigned_url=create_presigned_url(
AWS_S3_BUCKET_NAME, filename, expiration
),
filesize=get_file_size(AWS_S3_BUCKET_NAME, filename),
)
except HTTPError:
raise YtDlpStreamToS3Error("Could not connect to the media resource.")

Loading…
Cancel
Save