Skip to content

Commit 8874de3

Browse files
committed
feat: Cleaner redo of the archive wrappers.
1 parent b8a6d5a commit 8874de3

1 file changed

Lines changed: 39 additions & 30 deletions

File tree

‎src/tagstudio/previews/renderers/archive.py‎

Lines changed: 39 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,17 @@
22
# SPDX-License-Identifier: GPL-3.0-only
33

44

5-
import tarfile
6-
import zipfile
75
from io import BytesIO
86
from pathlib import Path
9-
from typing import Literal, override
7+
from tarfile import TarFile
8+
from typing import Self, override
9+
from zipfile import ZipFile
1010

11-
import py7zr
12-
import py7zr.io
13-
import rarfile
1411
import structlog
1512
from PIL.Image import Image
13+
from py7zr import SevenZipFile
14+
from py7zr.io import BytesIOFactory
15+
from rarfile import RarFile
1616

1717
from tagstudio.core.enums import Theme
1818
from tagstudio.core.media_types import MediaTypes
@@ -22,7 +22,7 @@
2222

2323
logger = structlog.get_logger(__name__)
2424

25-
type Archive = zipfile.ZipFile | rarfile.RarFile | SevenZipFile | TarFile
25+
type Archive = Rar | SevenZip | Tar | Zip
2626

2727

2828
class ArchivePreview(BasePreview):
@@ -57,41 +57,50 @@ def render(
5757
return archive_thumb(filepath)
5858

5959

60-
class SevenZipFile(py7zr.SevenZipFile):
61-
"""Wrapper around py7zr.SevenZipFile to mimic zipfile.ZipFile's API."""
60+
class Rar(RarFile):
61+
"""Wrapper around RarFile for a unified API."""
6262

63-
def __init__(self, filepath: Path, mode: Literal["r"]) -> None:
64-
super().__init__(filepath, mode)
63+
@classmethod
64+
def open_archive(cls, filepath: Path) -> Self:
65+
return cls(filepath, "r")
66+
67+
68+
class SevenZip(SevenZipFile):
69+
"""Wrapper around SevenZipFile for a unified API."""
70+
71+
@classmethod
72+
def open_archive(cls, filepath: Path) -> Self:
73+
return cls(filepath, "r")
6574

6675
def read(self, name: str) -> bytes:
6776
# SevenZipFile must be reset after every extraction
6877
# See https://py7zr.readthedocs.io/en/stable/api.html#py7zr.SevenZipFile.extract
6978
self.reset()
70-
factory = py7zr.io.BytesIOFactory(limit=10485760) # 10 MiB
79+
factory = BytesIOFactory(limit=10485760) # 10 MiB
7180
self.extract(targets=[name], factory=factory)
7281
return factory.get(name).read()
7382

7483

75-
class TarFile:
76-
"""Wrapper around tarfile.TarFile to mimic zipfile.ZipFile's API."""
84+
class Tar(TarFile):
85+
"""Wrapper around TarFile for a unified API."""
7786

78-
def __init__(self, filepath: Path, mode: Literal["r"]) -> None:
79-
self.tar: tarfile.TarFile
80-
self.filepath = filepath
81-
self.mode: Literal["r"] = mode
87+
@classmethod
88+
def open_archive(cls, filepath: Path) -> Self:
89+
return cls.open(filepath, "r")
8290

8391
def namelist(self) -> list[str]:
84-
return self.tar.getnames()
92+
return self.getnames()
8593

8694
def read(self, name: str) -> bytes:
87-
return unwrap(self.tar.extractfile(name)).read()
95+
return unwrap(self.extractfile(name)).read()
8896

89-
def __enter__(self) -> TarFile:
90-
self.tar = tarfile.open(name=self.filepath, mode=self.mode).__enter__()
91-
return self
9297

93-
def __exit__(self, *args) -> None: # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
94-
self.tar.__exit__(*args)
98+
class Zip(ZipFile):
99+
"""Wrapper around ZipFile for a unified API."""
100+
101+
@classmethod
102+
def open_archive(cls, filepath: Path) -> Self:
103+
return cls(filepath, "r")
95104

96105

97106
def open_archive(filepath: Path) -> Archive:
@@ -104,11 +113,11 @@ def open_archive(filepath: Path) -> Archive:
104113
Archive: The opened archive.
105114
"""
106115
exts = MediaTypes.get_equivalent_exts(MediaTypes.get_ext(filepath))
107-
archiver: type[Archive] = zipfile.ZipFile
116+
archiver: type[Archive] = Zip
108117
if exts & {".7z", ".cb7", ".s7z"}:
109-
archiver = SevenZipFile
118+
archiver = SevenZip
110119
elif exts & {".cbr", ".rar"}:
111-
archiver = rarfile.RarFile
120+
archiver = Rar
112121
elif exts & {
113122
".cbt",
114123
".tar",
@@ -119,8 +128,8 @@ def open_archive(filepath: Path) -> Archive:
119128
".tar.xz",
120129
".tar.zst",
121130
}:
122-
archiver = TarFile
123-
return archiver(filepath, "r")
131+
archiver = Tar
132+
return archiver.open_archive(filepath)
124133

125134

126135
def first_image_in_archive(archive: Archive) -> Image | None:

0 commit comments

Comments
 (0)