softwareweaver/MusicGen

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6"""Utility for reading some info from inside a zip file.7"""8 9import typing10import zipfile11 12from dataclasses import dataclass13from functools import lru_cache14from typing_extensions import Literal15 16 17DEFAULT_SIZE = 3218MODE = Literal['r', 'w', 'x', 'a']19 20 21@dataclass(order=True)22class PathInZip:23    """Hold a path of file within a zip file.24 25    Args:26        path (str): The convention is <path_to_zip>:<relative_path_inside_zip>.27            Let's assume there is a zip file /some/location/foo.zip28            and inside of it is a json file located at /data/file1.json,29            Then we expect path = "/some/location/foo.zip:/data/file1.json".30    """31 32    INFO_PATH_SEP = ':'33    zip_path: str34    file_path: str35 36    def __init__(self, path: str) -> None:37        split_path = path.split(self.INFO_PATH_SEP)38        assert len(split_path) == 239        self.zip_path, self.file_path = split_path40 41    @classmethod42    def from_paths(cls, zip_path: str, file_path: str):43        return cls(zip_path + cls.INFO_PATH_SEP + file_path)44 45    def __str__(self) -> str:46        return self.zip_path + self.INFO_PATH_SEP + self.file_path47 48 49def _open_zip(path: str, mode: MODE = 'r'):50    return zipfile.ZipFile(path, mode)51 52 53_cached_open_zip = lru_cache(DEFAULT_SIZE)(_open_zip)54 55 56def set_zip_cache_size(max_size: int):57    """Sets the maximal LRU caching for zip file opening.58 59    Args:60        max_size (int): the maximal LRU cache.61    """62    global _cached_open_zip63    _cached_open_zip = lru_cache(max_size)(_open_zip)64 65 66def open_file_in_zip(path_in_zip: PathInZip, mode: str = 'r') -> typing.IO:67    """Opens a file stored inside a zip and returns a file-like object.68 69    Args:70        path_in_zip (PathInZip): A PathInZip object representing the file to return a file-like object of.71        mode (str): The mode in which to open the file with.72    Returns:73        A file-like object for PathInZip.74    """75    zf = _cached_open_zip(path_in_zip.zip_path)76    return zf.open(path_in_zip.file_path)77