diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 2fd5676268e0c..1c1e117c36863 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -14,6 +14,8 @@ attrs==23.1.0 # via aiohttp black==23.7.0 # via -r requirements/requirements-dev.in +brotli==1.1.0 + # via -r requirements/requirements.in charset-normalizer==3.2.0 # via aiohttp click==8.1.6 diff --git a/requirements/requirements.in b/requirements/requirements.in index ee4ba4f3d739e..89bb017baa20a 100644 --- a/requirements/requirements.in +++ b/requirements/requirements.in @@ -1 +1,2 @@ aiohttp +brotli diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 26d32c7aa95c6..bb8bffa8070f7 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -12,6 +12,8 @@ async-timeout==4.0.2 # via aiohttp attrs==23.1.0 # via aiohttp +brotli==1.1.0 + # via -r requirements/requirements.in charset-normalizer==3.2.0 # via aiohttp frozenlist==1.4.0 diff --git a/src/file_manager.py b/src/file_manager.py index e06c42d7af4ce..2602e0515aa0f 100644 --- a/src/file_manager.py +++ b/src/file_manager.py @@ -111,6 +111,12 @@ def get_metadata_full_json_path(self) -> pathlib.Path: def get_metadata_compact_json_path(self) -> pathlib.Path: return self._get_metadata_dir() / "metadata-compact.json" + def get_metadata_full_json_br_path(self) -> pathlib.Path: + return self._get_metadata_dir() / "metadata-full.json.br" + + def get_metadata_compact_json_br_path(self) -> pathlib.Path: + return self._get_metadata_dir() / "metadata-compact.json.br" + def get_readme_path(self) -> pathlib.Path: return self._playlists_dir.parent / "README.md" diff --git a/src/file_updater.py b/src/file_updater.py index a7ba075259a16..942eab23b005e 100644 --- a/src/file_updater.py +++ b/src/file_updater.py @@ -4,7 +4,9 @@ import datetime import logging import pathlib -from typing import Dict, Optional, Set +from typing import Dict, Optional, Set, TypeVar + +import brotli from file_formatter import Formatter from file_manager import FileManager @@ -17,6 +19,9 @@ logger: logging.Logger = logging.getLogger(__name__) +T = TypeVar("T", str, bytes) + + class FileUpdater: @classmethod async def update_files( @@ -249,19 +254,27 @@ async def _update_files_impl( file_manager.ensure_no_unexpected_files() # Update all metadata files - metadata_full_content = Formatter.metadata_full_json(playlists) + "\n" - metadata_compact_content = Formatter.metadata_compact_json(playlists) + "\n" + metadata_full_json = Formatter.metadata_full_json(playlists) + metadata_compact_json = Formatter.metadata_compact_json(playlists) cls._maybe_update_file( path=file_manager.get_old_metadata_json_path(), - content=metadata_full_content, + content=metadata_full_json + "\n", ) cls._maybe_update_file( path=file_manager.get_metadata_full_json_path(), - content=metadata_full_content, + content=metadata_full_json + "\n", ) cls._maybe_update_file( path=file_manager.get_metadata_compact_json_path(), - content=metadata_compact_content, + content=metadata_compact_json + "\n", + ) + cls._maybe_update_file( + path=file_manager.get_metadata_full_json_br_path(), + content=brotli.compress(metadata_full_json.encode()), + ) + cls._maybe_update_file( + path=file_manager.get_metadata_compact_json_br_path(), + content=brotli.compress(metadata_compact_json.encode()), ) # Lastly, update README.md @@ -289,20 +302,39 @@ def _get_file_content_or_empty_string(cls, path: pathlib.Path) -> str: except FileNotFoundError: return "" + @classmethod + def _get_file_content_or_empty_bytes(cls, path: pathlib.Path) -> bytes: + try: + with open(path, "rb") as f: + return f.read() + except FileNotFoundError: + return b"" + @classmethod def _write_to_file_if_content_changed( - cls, prev_content: str, content: str, path: pathlib.Path + cls, prev_content: T, content: T, path: pathlib.Path ) -> None: if content == prev_content: logger.info(f" No changes to file: {path}") return logger.info(f" Writing updates to file: {path}") - with open(path, "w") as f: - f.write(content) + if isinstance(content, bytes): + with open(path, "wb") as f: + f.write(content) + elif isinstance(content, str): + with open(path, "w") as f: + f.write(content) + else: + raise RuntimeError(f"Invalid content type: {type(content)}") @classmethod - def _maybe_update_file(cls, path: pathlib.Path, content: str) -> None: - prev_content = cls._get_file_content_or_empty_string(path) + def _maybe_update_file(cls, path: pathlib.Path, content: T) -> None: + if isinstance(content, bytes): + prev_content = cls._get_file_content_or_empty_bytes(path) + elif isinstance(content, str): + prev_content = cls._get_file_content_or_empty_string(path) + else: + raise RuntimeError(f"Invalid content type: {type(content)}") cls._write_to_file_if_content_changed( prev_content=prev_content, content=content, diff --git a/src/tests/test_file_updater.py b/src/tests/test_file_updater.py index e0f8c65b58905..212cddf77c828 100644 --- a/src/tests/test_file_updater.py +++ b/src/tests/test_file_updater.py @@ -404,10 +404,35 @@ async def test_readme_and_metadata_json(self, mock_logger: Mock) -> None: """ ), ) + with open(metadata_dir / "metadata-compact.json", "r") as f: content = f.read() self.assertEqual(content, '{"a":"name_a","b":"name_b","c":" name_c "}\n') + with open(metadata_dir / "metadata-full.json.br", "rb") as f: + content = f.read() + self.assertEqual( + content, + ( + b"\x1b\x11\x03\x00\x1c\x07v,cz\xbe\xb1u'\xa6nK\xf5,$c\x1b\xdb-\x82" + b"\x1eQL&\x88H_\xea\xb0(L\xdd\xbc\xe7gYjc76\x8e\r\x9e>X\xf4\xc0\tQ" + b"\x17\x95n\x8b\x04\xf3W\x04\xe2\x8d;\xffH\xe0\xc6\x94z\x01\x9c\x1cu" + b"\xd4[Da\x03\xcd\xa76\xc9q\x04\x0ezG\xa5r\xd4u\xaf\x9eB\xb9S$f\xff" + b"\xe8\x9dy\x98\x81sz\xb9\xf9\x966D7\x1c\x1dL2Jl&4\x8dn\xc0\xd5\x8dB" + b"\xa5?\x9a\xf7\xf0\x0e&\x9a\x11?/\xc9\x87\xfc>C\xf4<\x81\x07\xb3j" + ), + ) + + with open(metadata_dir / "metadata-compact.json.br", "rb") as f: + content = f.read() + self.assertEqual( + content, + ( + b"\x1b)\x00\xf8\x1d\tv\xac\x89\xbb\xf348a\x08tc\xa9>7\xd9\x8fQC" + b"\x11C\xa4Xt:\x81EDqH\x15\xd0\xc0\x1e\x97\xe9\x82c\xa2\x14=" + ), + ) + async def test_success(self) -> None: # TODO pass