From fed0c2e4fe528b6fc5858b20faa2fe1b1166e53a Mon Sep 17 00:00:00 2001
From: Juan Olvera <juan@jolvera.dev>
Date: Mon, 5 Jun 2023 11:07:54 -0500
Subject: [PATCH] add filesize to output

---
 poetry.lock                   | 17 ++++++++++++++++-
 pyproject.toml                |  1 +
 yt_dlp_stream_to_s3/client.py | 20 +++++++++++++-------
 yt_dlp_stream_to_s3/main.py   | 15 ++++++++++-----
 4 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index a719ebc..e950fec 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -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"},
diff --git a/pyproject.toml b/pyproject.toml
index d8bf6e3..cf8a8e6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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]
diff --git a/yt_dlp_stream_to_s3/client.py b/yt_dlp_stream_to_s3/client.py
index ce9e405..1467bdf 100644
--- a/yt_dlp_stream_to_s3/client.py
+++ b/yt_dlp_stream_to_s3/client.py
@@ -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,
diff --git a/yt_dlp_stream_to_s3/main.py b/yt_dlp_stream_to_s3/main.py
index b26e7b7..e93465f 100644
--- a/yt_dlp_stream_to_s3/main.py
+++ b/yt_dlp_stream_to_s3/main.py
@@ -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.")