From 659eb048cc9cac73c46349eb29845bc5cd630f09 Mon Sep 17 00:00:00 2001
From: nergall2 <63109814+nergall2@users.noreply.github.com>
Date: Fri, 3 Apr 2020 18:42:21 +0300
Subject: [PATCH] enhancement-40175: Add remove() in ZipInfo

---
 Lib/zipfile.py | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index c3f814cc747e06..0699d40766b048 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -17,6 +17,7 @@
 import threading
 import time
 import contextlib
+from operator import attrgetter
 
 try:
     import zlib # We may need its compression method
@@ -1632,6 +1633,29 @@ def extractall(self, path=None, members=None, pwd=None):
 
         for zipinfo in members:
             self._extract_member(zipinfo, path, pwd)
+            
+    def remove(self, member):
+        """Remove a file from the archive. The archive must be open with mode 'a'"""
+
+        if self.mode != 'a':
+            raise RuntimeError("remove() requires mode 'a'")
+        if not self.fp:
+            raise ValueError(
+                "Attempt to write to ZIP archive that was already closed")
+        if self._writing:
+            raise ValueError(
+                "Can't write to ZIP archive while an open writing handle exists."
+            )
+
+        # Make sure we have an info object
+        if isinstance(member, ZipInfo):
+            # 'member' is already an info object
+            zinfo = member
+        else:
+            # get the info object
+            zinfo = self.getinfo(member)
+
+        return self._remove_member(zinfo)
 
     @classmethod
     def _sanitize_windows_name(cls, arcname, pathsep):
@@ -1689,6 +1713,52 @@ def _extract_member(self, member, targetpath, pwd):
             shutil.copyfileobj(source, target)
 
         return targetpath
+    
+    def _remove_member(self, member):
+        # get a sorted filelist by header offset, in case the dir order
+        # doesn't match the actual entry order
+        fp = self.fp
+        entry_offset = 0
+        filelist = sorted(self.filelist, key=attrgetter('header_offset'))
+        for i in range(len(filelist)):
+            info = filelist[i]
+            # find the target member
+            if info.header_offset < member.header_offset:
+                continue
+
+            # get the total size of the entry
+            entry_size = None
+            if i == len(filelist) - 1:
+                entry_size = self.start_dir - info.header_offset
+            else:
+                entry_size = filelist[i + 1].header_offset - info.header_offset
+
+            # found the member, set the entry offset
+            if member == info:
+                entry_offset = entry_size
+                continue
+
+            # Move entry
+            # read the actual entry data
+            fp.seek(info.header_offset)
+            entry_data = fp.read(entry_size)
+
+            # update the header
+            info.header_offset -= entry_offset
+
+            # write the entry to the new position
+            fp.seek(info.header_offset)
+            fp.write(entry_data)
+            fp.flush()
+
+        # update state
+        self.start_dir -= entry_offset
+        self.filelist.remove(member)
+        del self.NameToInfo[member.filename]
+        self._didModify = True
+
+        # seek to the start of the central dir
+        fp.seek(self.start_dir)
 
     def _writecheck(self, zinfo):
         """Check for errors before writing a file to the archive."""