Skip to content

Commit d008cb6

Browse files
yudilevi2danny0838
authored andcommitted
enhancement-40175: Add remove() in ZipInfo
(cherry picked from commit 659eb04 (PR #19358))
1 parent d706eb9 commit d008cb6

File tree

1 file changed

+71
-0
lines changed

1 file changed

+71
-0
lines changed

Lib/zipfile/__init__.py

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
import sys
1414
import threading
1515
import time
16+
import contextlib
17+
from operator import attrgetter
1618

1719
try:
1820
import zlib # We may need its compression method
@@ -1865,6 +1867,29 @@ def extractall(self, path=None, members=None, pwd=None):
18651867

18661868
for zipinfo in members:
18671869
self._extract_member(zipinfo, path, pwd)
1870+
1871+
def remove(self, member):
1872+
"""Remove a file from the archive. The archive must be open with mode 'a'"""
1873+
1874+
if self.mode != 'a':
1875+
raise RuntimeError("remove() requires mode 'a'")
1876+
if not self.fp:
1877+
raise ValueError(
1878+
"Attempt to write to ZIP archive that was already closed")
1879+
if self._writing:
1880+
raise ValueError(
1881+
"Can't write to ZIP archive while an open writing handle exists."
1882+
)
1883+
1884+
# Make sure we have an info object
1885+
if isinstance(member, ZipInfo):
1886+
# 'member' is already an info object
1887+
zinfo = member
1888+
else:
1889+
# get the info object
1890+
zinfo = self.getinfo(member)
1891+
1892+
return self._remove_member(zinfo)
18681893

18691894
@classmethod
18701895
def _sanitize_windows_name(cls, arcname, pathsep):
@@ -1929,6 +1954,52 @@ def _extract_member(self, member, targetpath, pwd):
19291954
shutil.copyfileobj(source, target)
19301955

19311956
return targetpath
1957+
1958+
def _remove_member(self, member):
1959+
# get a sorted filelist by header offset, in case the dir order
1960+
# doesn't match the actual entry order
1961+
fp = self.fp
1962+
entry_offset = 0
1963+
filelist = sorted(self.filelist, key=attrgetter('header_offset'))
1964+
for i in range(len(filelist)):
1965+
info = filelist[i]
1966+
# find the target member
1967+
if info.header_offset < member.header_offset:
1968+
continue
1969+
1970+
# get the total size of the entry
1971+
entry_size = None
1972+
if i == len(filelist) - 1:
1973+
entry_size = self.start_dir - info.header_offset
1974+
else:
1975+
entry_size = filelist[i + 1].header_offset - info.header_offset
1976+
1977+
# found the member, set the entry offset
1978+
if member == info:
1979+
entry_offset = entry_size
1980+
continue
1981+
1982+
# Move entry
1983+
# read the actual entry data
1984+
fp.seek(info.header_offset)
1985+
entry_data = fp.read(entry_size)
1986+
1987+
# update the header
1988+
info.header_offset -= entry_offset
1989+
1990+
# write the entry to the new position
1991+
fp.seek(info.header_offset)
1992+
fp.write(entry_data)
1993+
fp.flush()
1994+
1995+
# update state
1996+
self.start_dir -= entry_offset
1997+
self.filelist.remove(member)
1998+
del self.NameToInfo[member.filename]
1999+
self._didModify = True
2000+
2001+
# seek to the start of the central dir
2002+
fp.seek(self.start_dir)
19322003

19332004
def _writecheck(self, zinfo):
19342005
"""Check for errors before writing a file to the archive."""

0 commit comments

Comments
 (0)