From 02d47632cf91a858e0cb5fa8a40eb04e8b5654f5 Mon Sep 17 00:00:00 2001
From: ricterz <ricterzheng@gmail.com>
Date: Sun, 22 Sep 2024 15:07:53 +0800
Subject: [PATCH] fix bug of move-to-dir

---
 nhentai/command.py | 22 +++++++++++++++++-----
 nhentai/utils.py   | 40 +++++++++++++++++++++++-----------------
 2 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/nhentai/command.py b/nhentai/command.py
index 7feffcd..2e37b02 100644
--- a/nhentai/command.py
+++ b/nhentai/command.py
@@ -1,4 +1,6 @@
 # coding: utf-8
+import os
+import shutil
 import sys
 import signal
 import platform
@@ -12,7 +14,7 @@ from nhentai.downloader import Downloader
 from nhentai.logger import logger
 from nhentai.constant import BASE_URL
 from nhentai.utils import generate_html, generate_doc, generate_main_html, generate_metadata_file, \
-    paging, check_cookie, signal_handler, DB
+    paging, check_cookie, signal_handler, DB, move_to_folder
 
 
 def main():
@@ -104,12 +106,22 @@ def main():
                 generate_html(options.output_dir, doujinshi, template=constant.CONFIG['template'])
 
             if options.is_cbz:
-                generate_doc('cbz', options.output_dir, doujinshi, options.rm_origin_dir, options.move_to_folder,
-                             options.regenerate)
+                generate_doc('cbz', options.output_dir, doujinshi, options.rm_origin_dir)
 
             if options.is_pdf:
-                generate_doc('pdf', options.output_dir, doujinshi, options.rm_origin_dir, options.move_to_folder,
-                             options.regenerate)
+                generate_doc('pdf', options.output_dir, doujinshi, options.rm_origin_dir)
+
+            if options.move_to_folder:
+                if options.is_cbz:
+                    move_to_folder(options.output_dir, doujinshi, 'cbz')
+                if options.is_pdf:
+                    move_to_folder(options.output_dir, doujinshi, 'pdf')
+
+            if options.rm_origin_dir:
+                if options.move_to_folder:
+                    logger.critical('You specified both --move-to-folder and --rm-origin-dir options, '
+                                    'you will not get anything :(')
+                shutil.rmtree(os.path.join(options.output_dir, doujinshi.filename), ignore_errors=True)
 
         if options.main_viewer:
             generate_main_html(options.output_dir)
diff --git a/nhentai/utils.py b/nhentai/utils.py
index f22f8d7..73477ca 100644
--- a/nhentai/utils.py
+++ b/nhentai/utils.py
@@ -72,8 +72,8 @@ def parse_doujinshi_obj(
         doujinshi_obj=None,
         file_type: str = ''
 ) -> Tuple[str, str]:
-    filename = './doujinshi' + file_type
 
+    filename = f'./doujinshi.{file_type}'
     doujinshi_dir = os.path.join(output_dir, doujinshi_obj.filename)
     if doujinshi_obj is not None:
         _filename = f'{doujinshi_obj.filename}.{file_type}'
@@ -128,6 +128,27 @@ def generate_html(output_dir='.', doujinshi_obj=None, template='default'):
         logger.warning(f'Writing HTML Viewer failed ({e})')
 
 
+def move_to_folder(output_dir='.', doujinshi_obj=None, file_type=None):
+    if not file_type:
+        raise RuntimeError('no file_type specified')
+
+    doujinshi_dir, filename = parse_doujinshi_obj(output_dir, doujinshi_obj, file_type)
+
+    for fn in os.listdir(doujinshi_dir):
+        file_path = os.path.join(doujinshi_dir, fn)
+        _, ext = os.path.splitext(file_path)
+        if ext in ['.pdf', '.cbz']:
+            continue
+
+        if os.path.isfile(file_path):
+            try:
+                os.remove(file_path)
+            except Exception as e:
+                print(f"Error deleting file: {e}")
+
+    shutil.move(filename, os.path.join(doujinshi_dir, os.path.basename(filename)))
+
+
 def generate_main_html(output_dir='./'):
     """
     Generate a main html to show all the contains doujinshi.
@@ -185,8 +206,7 @@ def generate_main_html(output_dir='./'):
         logger.warning(f'Writing Main Viewer failed ({e})')
 
 
-def generate_doc(file_type='', output_dir='.', doujinshi_obj=None, rm_origin_dir=False,
-                 move_to_folder=False, regenerate=False):
+def generate_doc(file_type='', output_dir='.', doujinshi_obj=None, regenerate=False):
 
     doujinshi_dir, filename = parse_doujinshi_obj(output_dir, doujinshi_obj, file_type)
 
@@ -225,20 +245,6 @@ def generate_doc(file_type='', output_dir='.', doujinshi_obj=None, rm_origin_dir
         except ImportError:
             logger.error("Please install img2pdf package by using pip.")
 
-    if rm_origin_dir:
-        shutil.rmtree(doujinshi_dir, ignore_errors=True)
-
-    if move_to_folder:
-        for filename in os.listdir(doujinshi_dir):
-            file_path = os.path.join(doujinshi_dir, filename)
-            if os.path.isfile(file_path):
-                try:
-                    os.remove(file_path)
-                except Exception as e:
-                    print(f"Error deleting file: {e}")
-
-        shutil.move(filename, doujinshi_dir)
-
 
 def format_filename(s, length=MAX_FIELD_LENGTH, _truncate_only=False):
     """