22# SPDX-License-Identifier: GPL-3.0-only
33
44
5- import tarfile
6- import zipfile
75from io import BytesIO
86from 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
1411import structlog
1512from PIL .Image import Image
13+ from py7zr import SevenZipFile
14+ from py7zr .io import BytesIOFactory
15+ from rarfile import RarFile
1616
1717from tagstudio .core .enums import Theme
1818from tagstudio .core .media_types import MediaTypes
2222
2323logger = structlog .get_logger (__name__ )
2424
25- type Archive = zipfile . ZipFile | rarfile . RarFile | SevenZipFile | TarFile
25+ type Archive = Rar | SevenZip | Tar | Zip
2626
2727
2828class 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
97106def 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
126135def first_image_in_archive (archive : Archive ) -> Image | None :
0 commit comments