Skip to content

Commit f3450f1

Browse files
danny0838danny0838
committed
gh-51067: Add ZipInfo.remove() (#51067)
1 parent 30a306c commit f3450f1

File tree

5 files changed

+405
-0
lines changed

5 files changed

+405
-0
lines changed

‎Doc/library/zipfile.rst‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -492,6 +492,23 @@ ZipFile Objects
492492
.. versionadded:: 3.11
493493

494494

495+
.. method:: ZipFile.remove(zinfo_or_arcname)
496+
497+
Removes a member from the archive. *zinfo_or_arcname* is either the full
498+
path of the member, or a :class:`ZipInfo` instance.
499+
500+
The archive must be opened with mode ``'w'``, ``'x'`` or ``'a'``.
501+
502+
Calling :meth:`remove` on a closed ZipFile will raise a :exc:`ValueError`.
503+
504+
.. note::
505+
506+
Removing a member in an archive may involve a move of many internal data
507+
records, which can be I/O intensive for a large ZIP file.
508+
509+
.. versionadded:: 3.12
510+
511+
495512
The following data attributes are also available:
496513

497514
.. attribute:: ZipFile.filename

‎Lib/test/test_zipfile/test_core.py‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1147,6 +1147,239 @@ class LzmaWriterTests(AbstractWriterTests, unittest.TestCase):
11471147
compression = zipfile.ZIP_LZMA
11481148

11491149

1150+
class AbstractRemoveTests:
1151+
1152+
def _test_removing_indexes(self, test_files, indexes):
1153+
"""Test underlying _remove_members() for removing members at given
1154+
indexes."""
1155+
# calculate the expected results
1156+
expected_files = []
1157+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1158+
for i, (file, data) in enumerate(test_files):
1159+
if i not in indexes:
1160+
zh.writestr(file, data)
1161+
expected_files.append(file)
1162+
expected_size = os.path.getsize(TESTFN)
1163+
1164+
# prepare the test zip
1165+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1166+
for file, data in test_files:
1167+
zh.writestr(file, data)
1168+
1169+
# do the removal and check the result
1170+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1171+
members = {zh.infolist()[i] for i in indexes}
1172+
zh._remove_members(members)
1173+
1174+
# make sure internal caches have reflected the change
1175+
# and are consistent
1176+
self.assertEqual(zh.namelist(), expected_files)
1177+
for file, _ in test_files:
1178+
if file in zh.namelist():
1179+
self.assertEqual(zh.getinfo(file).filename, file)
1180+
else:
1181+
with self.assertRaises(KeyError):
1182+
zh.getinfo(file)
1183+
1184+
self.assertIsNone(zh.testzip())
1185+
self.assertEqual(os.path.getsize(TESTFN), expected_size)
1186+
1187+
def _test_removing_combinations(self, test_files, n=None):
1188+
"""Test underlying _remove_members() for removing random combinations
1189+
of members."""
1190+
ln = len(test_files)
1191+
if n is None:
1192+
# iterate n from 1 to all
1193+
for n in range(1, ln + 1):
1194+
for indexes in itertools.combinations(range(ln), n):
1195+
with self.subTest(remove=indexes):
1196+
self._test_removing_indexes(test_files, indexes)
1197+
else:
1198+
for indexes in itertools.combinations(range(ln), n):
1199+
with self.subTest(remove=indexes):
1200+
self._test_removing_indexes(test_files, indexes)
1201+
1202+
def test_basic(self):
1203+
# Test underlying _remove_members() for removing random combinations of members.
1204+
test_files = [
1205+
('file0.txt', b'Lorem ipsum dolor sit amet, consectetur adipiscing elit'),
1206+
('file1.txt', b'Duis aute irure dolor in reprehenderit in voluptate velit esse'),
1207+
('file2.txt', b'Sed ut perspiciatis unde omnis iste natus error sit voluptatem'),
1208+
]
1209+
1210+
self._test_removing_combinations(test_files)
1211+
1212+
def test_duplicated_arcname(self):
1213+
# Test underlying _remove_members() for removing any one of random duplicated members.
1214+
dupl_file = 'file.txt'
1215+
test_files = [
1216+
('file0.txt', b'Lorem ipsum dolor sit amet, consectetur adipiscing elit'),
1217+
('file1.txt', b'Duis aute irure dolor in reprehenderit in voluptate velit esse'),
1218+
('file2.txt', b'Sed ut perspiciatis unde omnis iste natus error sit voluptatem'),
1219+
]
1220+
1221+
ln = len(test_files)
1222+
for n in range(2, ln + 1):
1223+
for dups in itertools.combinations(range(ln), n):
1224+
files = []
1225+
for i, (file, data) in enumerate(test_files):
1226+
file_ = dupl_file if i in dups else file
1227+
files.append((file_, data))
1228+
1229+
for index in dups:
1230+
indexes = [index]
1231+
with self.subTest(dups=dups, indexes=indexes):
1232+
self._test_removing_indexes(files, indexes)
1233+
1234+
def test_non_physical(self):
1235+
# Test underlying _remove_members() for non-physical removing.
1236+
test_files = [
1237+
('file0.txt', b'Lorem ipsum dolor sit amet, consectetur adipiscing elit'),
1238+
('file1.txt', b'Duis aute irure dolor in reprehenderit in voluptate velit esse'),
1239+
('file2.txt', b'Sed ut perspiciatis unde omnis iste natus error sit voluptatem'),
1240+
]
1241+
1242+
ln = len(test_files)
1243+
for n in range(1, ln + 1):
1244+
for indexes in itertools.combinations(range(ln), n):
1245+
with self.subTest(remove=indexes):
1246+
# prepare the test zip
1247+
expected = {}
1248+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1249+
for i, (file, data) in enumerate(test_files):
1250+
zh.writestr(file, data)
1251+
if i not in indexes:
1252+
expected[file] = zh.getinfo(file).header_offset
1253+
1254+
# do the removal and check the result
1255+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1256+
members = {zh.infolist()[i] for i in indexes}
1257+
zh._remove_members(members, remove_physical=False)
1258+
self.assertEqual(zh.namelist(), list(expected))
1259+
for file, offset in expected.items():
1260+
self.assertEqual(zh.getinfo(file).header_offset, offset)
1261+
self.assertIsNone(zh.testzip())
1262+
1263+
def test_verify(self):
1264+
# Test if params are passed to underlying _remove_members() correctly,
1265+
# or never passed if conditions not met.
1266+
file0 = 'file0.txt'
1267+
file = 'datafile.txt'
1268+
data = b'Sed ut perspiciatis unde omnis iste natus error sit voluptatem'
1269+
1270+
# closed: error and do nothing
1271+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1272+
zh.writestr(file, data)
1273+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1274+
zh.close()
1275+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1276+
with self.assertRaises(ValueError):
1277+
zh.remove(file)
1278+
mock_fn.assert_not_called()
1279+
1280+
# writing: error and do nothing
1281+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1282+
zh.writestr(file, data)
1283+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1284+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1285+
with zh.open(file0, 'w') as fh:
1286+
with self.assertRaises(ValueError):
1287+
zh.remove(file)
1288+
mock_fn.assert_not_called()
1289+
1290+
# mode 'r': error and do nothing
1291+
with zipfile.ZipFile(TESTFN, 'r') as zh:
1292+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1293+
with self.assertRaises(ValueError):
1294+
zh.remove(file)
1295+
mock_fn.assert_not_called()
1296+
1297+
# mode 'a': the most general use case
1298+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1299+
zh.writestr(file, data)
1300+
# -- remove with arcname
1301+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1302+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1303+
zh.remove(file)
1304+
mock_fn.assert_called_once_with({zh.getinfo(file)})
1305+
# -- remove with zinfo
1306+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1307+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1308+
zinfo = zh.getinfo(file)
1309+
zh.remove(zinfo)
1310+
mock_fn.assert_called_once_with({zinfo})
1311+
# -- remove with nonexist arcname
1312+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1313+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1314+
with self.assertRaises(KeyError):
1315+
zh.remove('nonexist.file')
1316+
mock_fn.assert_not_called()
1317+
# -- remove with nonexist zinfo (even if same name)
1318+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1319+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1320+
zinfo = zipfile.ZipInfo(file)
1321+
with self.assertRaises(KeyError):
1322+
zh.remove(zinfo)
1323+
mock_fn.assert_not_called()
1324+
1325+
# mode 'w': like 'a'; allows removing a just written member
1326+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1327+
zh.writestr(file, data)
1328+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1329+
zh.remove(file)
1330+
mock_fn.assert_called_once_with({zh.getinfo(file)})
1331+
1332+
# mode 'x': like 'w'
1333+
os.remove(TESTFN)
1334+
with zipfile.ZipFile(TESTFN, 'x') as zh:
1335+
zh.writestr(file, data)
1336+
with mock.patch('zipfile.ZipFile._remove_members') as mock_fn:
1337+
zh.remove(file)
1338+
mock_fn.assert_called_once_with({zh.getinfo(file)})
1339+
1340+
def test_zip64(self):
1341+
# Test if members use zip64.
1342+
file = 'datafile.txt'
1343+
file1 = 'pre.txt'
1344+
file2 = 'post.txt'
1345+
data = b'Sed ut perspiciatis unde omnis iste natus error sit voluptatem'
1346+
data1 = b'Lorem ipsum dolor sit amet, consectetur adipiscing elit'
1347+
data2 = b'Duis aute irure dolor in reprehenderit in voluptate velit esse'
1348+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1349+
with zh.open(file1, 'w', force_zip64=True) as fh:
1350+
fh.write(data1)
1351+
with zh.open(file2, 'w', force_zip64=True) as fh:
1352+
fh.write(data2)
1353+
expected_size = os.path.getsize(TESTFN)
1354+
1355+
with zipfile.ZipFile(TESTFN, 'w') as zh:
1356+
with zh.open(file1, 'w', force_zip64=True) as fh:
1357+
fh.write(data1)
1358+
with zh.open(file, 'w', force_zip64=True) as fh:
1359+
fh.write(data)
1360+
with zh.open(file2, 'w', force_zip64=True) as fh:
1361+
fh.write(data2)
1362+
with zipfile.ZipFile(TESTFN, 'a') as zh:
1363+
zh.remove(file)
1364+
self.assertIsNone(zh.testzip())
1365+
self.assertEqual(os.path.getsize(TESTFN), expected_size)
1366+
1367+
class StoredRemoveTests(AbstractRemoveTests, unittest.TestCase):
1368+
compression = zipfile.ZIP_STORED
1369+
1370+
@requires_zlib()
1371+
class DeflateRemoveTests(AbstractRemoveTests, unittest.TestCase):
1372+
compression = zipfile.ZIP_DEFLATED
1373+
1374+
@requires_bz2()
1375+
class Bzip2RemoveTests(AbstractRemoveTests, unittest.TestCase):
1376+
compression = zipfile.ZIP_BZIP2
1377+
1378+
@requires_lzma()
1379+
class LzmaRemoveTests(AbstractRemoveTests, unittest.TestCase):
1380+
compression = zipfile.ZIP_LZMA
1381+
1382+
11501383
class PyZipFileTests(unittest.TestCase):
11511384
def assertCompiledIn(self, name, namelist):
11521385
if name + 'o' not in namelist:

‎Lib/test/test_zipfile64.py‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,69 @@ def tearDown(self):
8787
os_helper.unlink(TESTFN2)
8888

8989

90+
class TestRemove(unittest.TestCase):
91+
def setUp(self):
92+
# Create test data.
93+
line_gen = ("Test of zipfile line %d." % i for i in range(1000000))
94+
self.data = '\n'.join(line_gen).encode('ascii')
95+
96+
def _write_large_file(self, fh):
97+
# It will contain enough copies of self.data to reach about 8 GiB.
98+
filecount = 8*1024**3 // len(self.data)
99+
100+
next_time = time.monotonic() + _PRINT_WORKING_MSG_INTERVAL
101+
for num in range(filecount):
102+
fh.write(self.data)
103+
# Print still working message since this test can be really slow
104+
if next_time <= time.monotonic():
105+
next_time = time.monotonic() + _PRINT_WORKING_MSG_INTERVAL
106+
print((
107+
' writing %d of %d, be patient...' %
108+
(num, filecount)), file=sys.__stdout__)
109+
sys.__stdout__.flush()
110+
111+
def test_remove_large_file(self):
112+
# Try the temp file. If we do TESTFN2, then it hogs
113+
# gigabytes of disk space for the duration of the test.
114+
with TemporaryFile() as f:
115+
self._test_remove_large_file(f)
116+
self.assertFalse(f.closed)
117+
118+
def _test_remove_large_file(self, f):
119+
file = 'datafile.txt'
120+
file1 = 'dummy.txt'
121+
data = b'Sed ut perspiciatis unde omnis iste natus error sit voluptatem'
122+
with zipfile.ZipFile(f, 'w') as zh:
123+
with zh.open(file1, 'w', force_zip64=True) as fh:
124+
self._write_large_file(fh)
125+
zh.writestr(file, data)
126+
127+
with zipfile.ZipFile(f, 'a') as zh:
128+
zh.remove(file1)
129+
self.assertIsNone(zh.testzip())
130+
131+
def test_remove_before_large_file(self):
132+
# Try the temp file. If we do TESTFN2, then it hogs
133+
# gigabytes of disk space for the duration of the test.
134+
with TemporaryFile() as f:
135+
self._test_remove_before_large_file(f)
136+
self.assertFalse(f.closed)
137+
138+
def _test_remove_before_large_file(self, f):
139+
file = 'datafile.txt'
140+
file1 = 'dummy.txt'
141+
data = b'Sed ut perspiciatis unde omnis iste natus error sit voluptatem'
142+
with zipfile.ZipFile(f, 'w') as zh:
143+
zh.writestr(file, data)
144+
with zh.open(file1, 'w', force_zip64=True) as fh:
145+
self._write_large_file(fh)
146+
expected_size = zh.getinfo(file1).file_size
147+
148+
with zipfile.ZipFile(f, 'a') as zh:
149+
zh.remove(file)
150+
self.assertIsNone(zh.testzip())
151+
152+
90153
class OtherTests(unittest.TestCase):
91154
def testMoreThan64kFiles(self):
92155
# This test checks that more than 64k files can be added to an archive,

0 commit comments

Comments
 (0)