|
17 | 17 | import threading
|
18 | 18 | import time
|
19 | 19 | import contextlib
|
| 20 | +from operator import attrgetter |
20 | 21 |
|
21 | 22 | try:
|
22 | 23 | import zlib # We may need its compression method
|
@@ -1632,6 +1633,29 @@ def extractall(self, path=None, members=None, pwd=None):
|
1632 | 1633 |
|
1633 | 1634 | for zipinfo in members:
|
1634 | 1635 | self._extract_member(zipinfo, path, pwd)
|
| 1636 | + |
| 1637 | + def remove(self, member): |
| 1638 | + """Remove a file from the archive. The archive must be open with mode 'a'""" |
| 1639 | + |
| 1640 | + if self.mode != 'a': |
| 1641 | + raise RuntimeError("remove() requires mode 'a'") |
| 1642 | + if not self.fp: |
| 1643 | + raise ValueError( |
| 1644 | + "Attempt to write to ZIP archive that was already closed") |
| 1645 | + if self._writing: |
| 1646 | + raise ValueError( |
| 1647 | + "Can't write to ZIP archive while an open writing handle exists." |
| 1648 | + ) |
| 1649 | + |
| 1650 | + # Make sure we have an info object |
| 1651 | + if isinstance(member, ZipInfo): |
| 1652 | + # 'member' is already an info object |
| 1653 | + zinfo = member |
| 1654 | + else: |
| 1655 | + # get the info object |
| 1656 | + zinfo = self.getinfo(member) |
| 1657 | + |
| 1658 | + return self._remove_member(zinfo) |
1635 | 1659 |
|
1636 | 1660 | @classmethod
|
1637 | 1661 | def _sanitize_windows_name(cls, arcname, pathsep):
|
@@ -1689,6 +1713,52 @@ def _extract_member(self, member, targetpath, pwd):
|
1689 | 1713 | shutil.copyfileobj(source, target)
|
1690 | 1714 |
|
1691 | 1715 | return targetpath
|
| 1716 | + |
| 1717 | + def _remove_member(self, member): |
| 1718 | + # get a sorted filelist by header offset, in case the dir order |
| 1719 | + # doesn't match the actual entry order |
| 1720 | + fp = self.fp |
| 1721 | + entry_offset = 0 |
| 1722 | + filelist = sorted(self.filelist, key=attrgetter('header_offset')) |
| 1723 | + for i in range(len(filelist)): |
| 1724 | + info = filelist[i] |
| 1725 | + # find the target member |
| 1726 | + if info.header_offset < member.header_offset: |
| 1727 | + continue |
| 1728 | + |
| 1729 | + # get the total size of the entry |
| 1730 | + entry_size = None |
| 1731 | + if i == len(filelist) - 1: |
| 1732 | + entry_size = self.start_dir - info.header_offset |
| 1733 | + else: |
| 1734 | + entry_size = filelist[i + 1].header_offset - info.header_offset |
| 1735 | + |
| 1736 | + # found the member, set the entry offset |
| 1737 | + if member == info: |
| 1738 | + entry_offset = entry_size |
| 1739 | + continue |
| 1740 | + |
| 1741 | + # Move entry |
| 1742 | + # read the actual entry data |
| 1743 | + fp.seek(info.header_offset) |
| 1744 | + entry_data = fp.read(entry_size) |
| 1745 | + |
| 1746 | + # update the header |
| 1747 | + info.header_offset -= entry_offset |
| 1748 | + |
| 1749 | + # write the entry to the new position |
| 1750 | + fp.seek(info.header_offset) |
| 1751 | + fp.write(entry_data) |
| 1752 | + fp.flush() |
| 1753 | + |
| 1754 | + # update state |
| 1755 | + self.start_dir -= entry_offset |
| 1756 | + self.filelist.remove(member) |
| 1757 | + del self.NameToInfo[member.filename] |
| 1758 | + self._didModify = True |
| 1759 | + |
| 1760 | + # seek to the start of the central dir |
| 1761 | + fp.seek(self.start_dir) |
1692 | 1762 |
|
1693 | 1763 | def _writecheck(self, zinfo):
|
1694 | 1764 | """Check for errors before writing a file to the archive."""
|
|
0 commit comments