fix file renaming after task completion

Author: nasbdh9

xeHentai/task.py

@@ -258,30 +258,14 @@ def save_file(self, imgurl, redirect_url, binary_iter):
 
         self._f_lock.acquire()
         try:
-            try:
-                shutil.move(fn_tmp, fn)
-            except WindowsError as ex:
-                # file is used by another process
-                # do a copy and delete, WindowsError[32]
-                if ex.errno == 13:
-                    shutil.copy(fn_tmp, fn)
-                    try:
-                        os.unlink(fn_tmp)
-                    except:
-                        pass
-                else:
-                    raise ex
+            # Move the temp file to the final destination
+            shutil.move(fn_tmp, fn)
             self.set_fid_finished(fid)
-            if fid in self.duplicate_map:
-                for fid_rep in self.duplicate_map[fid]:
-                    # if a file download is interrupted, it will appear in self.filehash_map as well
-                    if fid_rep == fid:
-                        continue
-                    fn_rep = os.path.join(fpath, self.get_fidpad(fid_rep))
-                    shutil.copyfile(fn, fn_rep)
-                    self.set_fid_finished(fid_rep)
-                self.logger.debug("#%s is copied from #%s in save_file" % (fid_rep, fid))
-                del self.duplicate_map[fid]
+
+            # Store the actual downloaded file name and its extension
+            actual_ext = os.path.splitext(fname)[1] or '.jpg'  # Default to .jpg if no extension
+            self.renamed_map[fid] = fname  # Store the complete filename for renaming later
+
         except Exception as ex:
             self._f_lock.release()
             raise ex
@@ -331,60 +315,37 @@ def rename_fname(self):
         tmppath = os.path.join(fpath, RENAME_TMPDIR)
         cnt = 0
         error_list = []
-        # we need to track renamed fid's to decide 
-        # whether to rename into a temp filename or add (1)
-        # only need it when rename_ori = True
-        done_list = set()
         for fid in list(self.renamed_map.keys()):
-            fname = self.renamed_map[fid]
-            original_ext = os.path.splitext(fname)[1]
-            if original_ext== "":
-                original_ext = os.path.splitext(fname)[0]
-            # if we don't need to rename to original name and file type matches
-            if not self.config['rename_ori'] and original_ext.lower() == '.jpg':
-                continue
-            fname_ori = os.path.join(fpath, self.get_fidpad(fid)) # id   
+            # Check if rename_ori is enabled
             if self.config['rename_ori']:
-                if os.path.exists(os.path.join(tmppath, self.get_fidpad(fid))):
-                    # if we previously put it into a temporary folder, we need to change fname_ori
-                    fname_ori = os.path.join(tmppath, self.get_fidpad(fid))
-                fname_to = os.path.join(fpath, util.legalpath(fname))
+                # Use the stored full filename
+                fname_new = os.path.join(fpath, util.legalpath(self.renamed_map[fid]))
+                fname_ori = os.path.join(fpath, self.get_fidpad(fid))  # Original filename with ID
             else:
                 # Q: Why we don't just use id.ext when saving files instead of using
                 #   id.jpg?
                 # A: If former task doesn't download all files, a new task with same gallery
                 #   will have zero knowledge about file type before scanning all per page,
                 #   thus can't determine if this id is downloaded, because file type is not
                 #   necessarily .jpg
-                fname_to = os.path.join(fpath, self.get_fidpad(fid, original_ext[1:]))
-            while fname_ori != fname_to:
+                #
+                # Use the previously stored extension
+                original_ext = os.path.splitext(self.renamed_map[fid])[1]  # Get the extension from the stored filename
+                fname_ori = os.path.join(fpath, self.get_fidpad(fid))  # ID-based filename
+                fname_new = os.path.join(fpath, self.get_fidpad(fid, original_ext[1:]))  # Use the stored extension
+
+            while fname_ori != fname_new:
                 if os.path.exists(fname_ori):
-                    while os.path.exists(fname_to):
-                        _base, _ext = os.path.splitext(fname_to)
-                        _ = re.findall("\((\d+)\)$", _base)
-                        if self.config['rename_ori'] and fname_to not in done_list:
-                            # if our auto numbering conflicts with original naming
-                            # we move it into a temporary folder
-                            # It's safe since this file is same with one of our auto numbering filename,
-                            # it could never be conflicted with other files in tmppath
-                            if not os.path.exists(tmppath):
-                                os.mkdir(tmppath)
-                            os.rename(fname_to, os.path.join(tmppath, os.path.split(fname_to)[1]))
-                            break
-                        if _ :# if ...(1) exists, use ...(2)
-                            print(safestr(_base))
-                            _base = re.sub("\((\d+)\)$", _base, lambda x:"(%d)" % (int(x.group(1)) + 1))
-                        else:
-                            _base = "%s(1)" % _base
-                        fname_to = "".join((_base, _ext))
+                    while os.path.exists(fname_new):
+                        _base, _ext = os.path.splitext(fname_new)
+                        _base = "%s(1)" % _base
+                        fname_new = "".join((_base, original_ext))
                 try:
-                    os.rename(fname_ori, fname_to)
-                    self.renamed_map[str(fid)] = os.path.split(fname_to)[1]
+                    os.rename(fname_ori, fname_new)
+                    self.renamed_map[str(fid)] = os.path.split(fname_new)[1]
                 except Exception as ex:
-                    error_list.append((os.path.split(fname_ori)[1], os.path.split(fname_to)[1], str(ex)))
+                    error_list.append((os.path.split(fname_ori)[1], os.path.split(fname_new)[1], str(ex)))
                     break
-                if self.config['rename_ori']:
-                    done_list.add(fname_to)
                 break
             cnt += 1
         if cnt == self.meta['total']: