|
13 | 13 | import sys |
14 | 14 | import threading |
15 | 15 | import time |
| 16 | +import contextlib |
| 17 | +from operator import attrgetter |
16 | 18 |
|
17 | 19 | try: |
18 | 20 | import zlib # We may need its compression method |
@@ -1865,6 +1867,29 @@ def extractall(self, path=None, members=None, pwd=None): |
1865 | 1867 |
|
1866 | 1868 | for zipinfo in members: |
1867 | 1869 | 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) |
1868 | 1893 |
|
1869 | 1894 | @classmethod |
1870 | 1895 | def _sanitize_windows_name(cls, arcname, pathsep): |
@@ -1929,6 +1954,52 @@ def _extract_member(self, member, targetpath, pwd): |
1929 | 1954 | shutil.copyfileobj(source, target) |
1930 | 1955 |
|
1931 | 1956 | 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) |
1932 | 2003 |
|
1933 | 2004 | def _writecheck(self, zinfo): |
1934 | 2005 | """Check for errors before writing a file to the archive.""" |
|
0 commit comments