Spaces:
Running on Zero
Running on Zero
File size: 7,138 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | """Zip data backend.
This backend works with filepaths pointing to valid Zip files. We assume that
the given Zip file contains the whole dataset associated to this backend.
"""
from __future__ import annotations
import os
import zipfile
from typing import Literal
from zipfile import ZipFile
from .base import DataBackend
class ZipBackend(DataBackend):
"""Backend for loading data from Zip files.
This backend works with filepaths pointing to valid Zip files. We assume
that the given Zip file contains the whole dataset associated to this
backend.
"""
def __init__(self) -> None:
"""Creates an instance of the class."""
super().__init__()
self.db_cache: dict[str, tuple[ZipFile, str]] = {}
@staticmethod
def _get_zip_path(
filepath: str, allow_omitted_ext: bool = True
) -> tuple[str, list[str]]:
"""Get .zip path and keys from filepath.
Args:
filepath (str): The filepath to retrieve the data from.
Should have the following format: 'path/to/file.zip/key1/key2'
allow_omitted_ext (bool, optional): Whether to allow omitted
extension, in which case the backend will try to append
'.zip' to the filepath. Defaults to True.
Returns:
tuple[str, list[str]]: The .hdf5 path and the keys to retrieve.
Examples:
>>> _get_zip_path("path/to/file.zip/key1/key2")
("path/to/file.zip", ["key2", "key1"])
>>> _get_zip_path("path/to/file/key1/key2", True)
("path/to/file.zip", ["key2", "key1"]) # if file.hdf5 exists and
# is a valid hdf5 file
"""
filepath_as_list = filepath.split("/")
keys = []
while True:
if filepath.endswith(".zip") or filepath == "":
break
if allow_omitted_ext and zipfile.is_zipfile(filepath + ".zip"):
filepath = filepath + ".zip"
break
keys.append(filepath_as_list.pop())
filepath = "/".join(filepath_as_list)
return filepath, keys
def exists(self, filepath: str) -> bool:
"""Check if filepath exists.
Args:
filepath (str): Path to file.
Returns:
bool: True if file exists, False otherwise.
"""
zip_path, keys = self._get_zip_path(filepath)
if not os.path.exists(zip_path):
return False
file = self._get_client(zip_path, "r")
url = "/".join(reversed(keys))
return url in file.namelist()
def set(
self, filepath: str, content: bytes, mode: Literal["w", "a"] = "w"
) -> None:
"""Write the file content to the zip file.
Args:
filepath: path/to/file.zip/key1/key2/key3
content: Bytes to be written to entry key3 within group key2
within another group key1, for example.
mode: Mode to open the file in. "w" for writing a file, "a" for
appending to existing file.
Raises:
ValueError: If filepath is not a valid .zip file
NotImplementedError: If the method is not implemented.
"""
if ".zip" not in filepath:
raise ValueError(f"{filepath} not a valid .zip filepath!")
zip_path, keys = self._get_zip_path(filepath)
zip_file = self._get_client(zip_path, mode)
url = "/".join(reversed(keys))
zip_file.writestr(url, content)
def _get_client(
self, zip_path: str, mode: Literal["r", "w", "a", "x"]
) -> ZipFile:
"""Get Zip client from path.
Args:
zip_path (str): Path to Zip file.
mode (str): Mode to open the file in.
Returns:
ZipFile: the hdf5 file.
"""
assert len(mode) == 1, "Mode must be a single character for zip file."
if zip_path not in self.db_cache:
os.makedirs(os.path.dirname(zip_path), exist_ok=True)
client = ZipFile(zip_path, mode)
self.db_cache[zip_path] = (client, mode)
else:
client, current_mode = self.db_cache[zip_path]
if current_mode != mode:
client.close()
client = ZipFile( # pylint:disable=consider-using-with
zip_path, mode
)
self.db_cache[zip_path] = (client, mode)
return client
def get(self, filepath: str) -> bytes:
"""Get values according to the filepath as bytes.
Args:
filepath (str): The path to the file. It consists of an Zip path
together with the relative path inside it, e.g.: "/path/to/
file.zip/key/subkey/data". If no .zip given inside filepath,
the function will search for the first .zip file present in
the path, i.e. "/path/to/file/key/subkey/data" will also /key/
subkey/data from /path/to/file.zip.
Raises:
ZipFileNotFoundError: If no suitable file exists.
OSError: If the file cannot be opened.
ValueError: If key not found inside zip file.
Returns:
bytes: The file content in bytes
"""
zip_path, keys = self._get_zip_path(filepath)
if not os.path.exists(zip_path):
raise FileNotFoundError(
f"Corresponding zip file not found:" f" {filepath}"
)
zip_file = self._get_client(zip_path, "r")
url = "/".join(reversed(keys))
try:
with zip_file.open(url) as zf:
content = zf.read()
except KeyError as e:
raise ValueError(f"Value '{url}' not found in {zip_path}!") from e
return bytes(content)
def listdir(self, filepath: str) -> list[str]:
"""List all files in the given directory.
Args:
filepath (str): The path to the directory.
Returns:
list[str]: List of all files in the given directory.
"""
zip_path, keys = self._get_zip_path(filepath)
zip_file = self._get_client(zip_path, "r")
url = "/".join(reversed(keys))
files = [
os.path.basename(key)
for key in zip_file.namelist()
if key.startswith(url) and os.path.basename(key) != ""
]
return sorted(files)
def isfile(self, filepath: str) -> bool:
"""Check if filepath is a file.
Args:
filepath (str): Path to file.
Returns:
bool: True if file exists, False otherwise.
"""
zip_path, keys = self._get_zip_path(filepath)
if not os.path.exists(zip_path):
return False
zip_file = self._get_client(zip_path, "r")
url = "/".join(reversed(keys))
return url in zip_file.namelist()
def close(self) -> None:
"""Close all opened Zip files."""
for client, _ in self.db_cache.values():
client.close()
self.db_cache = {}
|