Skip to content

Commit 3a2147e

Browse files
committed
Introduce _Extra.update()
Introduce a helper `.update()` that simplifies the common operation of fields updating of extra data bytes, by automatically stripping existing subfields with identical IDs and inserts multiple fields in one step.
1 parent b1ef21a commit 3a2147e

1 file changed

Lines changed: 15 additions & 3 deletions

File tree

Lib/zipfile/__init__.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,18 @@ def strip(cls, data, xids):
232232
if ex.id not in xids
233233
)
234234

235+
@classmethod
236+
def update(cls, data, extra):
237+
"""Insert fields from extra and strip duplicates."""
238+
extras = {
239+
ex.id: ex
240+
for ex in cls.split(extra, True)
241+
if ex.id is not None
242+
}
243+
# New fields first since data may have a corrupted tail that renders
244+
# following fields inaccessible.
245+
return b''.join(extras.values()) + cls.strip(data, extras)
246+
235247

236248
def _check_zipfile(fp):
237249
try:
@@ -2663,10 +2675,10 @@ def _write_end_record(self):
26632675
min_version = 0
26642676
if extra:
26652677
# Prepend a ZIP64 field to the extra's
2666-
extra_data = _Extra.strip(extra_data, (1,))
2667-
extra_data = struct.pack(
2678+
extra_data = _Extra.update(extra_data, struct.pack(
26682679
'<HH' + 'Q'*len(extra),
2669-
1, 8*len(extra), *extra) + extra_data
2680+
1, 8*len(extra), *extra,
2681+
))
26702682

26712683
min_version = ZIP64_VERSION
26722684

0 commit comments

Comments
 (0)