1
1
from pathlib import Path
2
2
import yaml
3
- from PySide6 .QtWidgets import (QFileDialog , QDialog , QVBoxLayout , QTextEdit ,
4
- QPushButton , QHBoxLayout , QMessageBox )
5
- from create_symlinks import create_symlinks_parallel
6
-
7
- ALLOWED_EXTENSIONS = {'.pdf' , '.docx' , '.epub' , '.txt' , '.enex' , '.eml' , '.msg' , '.csv' , '.xls' , '.xlsx' ,
8
- '.rtf' , '.odt' , '.png' , '.jpg' , '.jpeg' , '.bmp' , '.gif' , '.tif' , '.tiff' , '.html' ,
9
- '.htm' , '.md' , '.doc' }
3
+ from PySide6 .QtCore import QThread , Signal , Qt , QElapsedTimer
4
+ from PySide6 .QtWidgets import (QFileDialog , QDialog , QVBoxLayout , QTextEdit , QPushButton , QHBoxLayout , QMessageBox , QProgressDialog , QApplication , QFileSystemModel )
5
+ from multiprocessing import Pool , cpu_count
6
+ from create_symlinks import _create_single_symlink
7
+ ALLOWED_EXTENSIONS = {'.pdf' , '.docx' , '.epub' , '.txt' , '.enex' , '.eml' , '.msg' , '.csv' , '.xls' , '.xlsx' , '.rtf' , '.odt' , '.png' , '.jpg' , '.jpeg' , '.bmp' , '.gif' , '.tif' , '.tiff' , '.html' , '.htm' , '.md' , '.doc' }
10
8
DOCS_FOLDER = "Docs_for_DB"
11
9
CONFIG_FILE = "config.yaml"
12
-
10
+ class SymlinkWorker (QThread ):
11
+ progress = Signal (int )
12
+ finished = Signal (int , list )
13
+ def __init__ (self , source , target_dir , parent = None ):
14
+ super ().__init__ (parent )
15
+ self .source = source
16
+ self .target_dir = Path (target_dir )
17
+ def run (self ):
18
+ if isinstance (self .source , (str , Path )):
19
+ dir_path = Path (self .source )
20
+ files = [str (p ) for p in dir_path .iterdir () if p .is_file () and p .suffix .lower () in ALLOWED_EXTENSIONS ]
21
+ else :
22
+ files = list (self .source )
23
+ total = len (files )
24
+ made = 0
25
+ errors = []
26
+ last_pct = - 1
27
+ timer = QElapsedTimer ()
28
+ timer .start ()
29
+ step = max (1 , total // 100 ) if total else 1
30
+ if total > 1000 :
31
+ processes = min ((total // 10000 ) + 1 , cpu_count ())
32
+ file_args = [(f , str (self .target_dir )) for f in files ]
33
+ with Pool (processes = processes ) as pool :
34
+ for i , (ok , err ) in enumerate (pool .imap_unordered (_create_single_symlink , file_args ), 1 ):
35
+ if ok :
36
+ made += 1
37
+ if err :
38
+ errors .append (err )
39
+ if i % step == 0 or i == total :
40
+ pct = int (i * 100 / total ) if total else 100
41
+ if pct != last_pct and timer .elapsed () > 500 :
42
+ self .progress .emit (pct )
43
+ last_pct = pct
44
+ timer .restart ()
45
+ else :
46
+ for f in files :
47
+ if self .isInterruptionRequested ():
48
+ break
49
+ ok , err = _create_single_symlink ((f , str (self .target_dir )))
50
+ if ok :
51
+ made += 1
52
+ if err :
53
+ errors .append (err )
54
+ if made % step == 0 or made == total :
55
+ pct = int (made * 100 / total ) if total else 100
56
+ if pct != last_pct and timer .elapsed () > 500 :
57
+ self .progress .emit (pct )
58
+ last_pct = pct
59
+ timer .restart ()
60
+ self .finished .emit (made , errors )
13
61
def choose_documents_directory ():
14
62
current_dir = Path (__file__ ).parent .resolve ()
15
63
target_dir = current_dir / DOCS_FOLDER
16
64
target_dir .mkdir (parents = True , exist_ok = True )
17
-
18
65
msg_box = QMessageBox ()
19
66
msg_box .setWindowTitle ("Selection Type" )
20
67
msg_box .setText ("Would you like to select a directory or individual files?" )
21
68
dir_button = msg_box .addButton ("Select Directory" , QMessageBox .ActionRole )
22
69
files_button = msg_box .addButton ("Select Files" , QMessageBox .ActionRole )
23
70
cancel_button = msg_box .addButton ("Cancel" , QMessageBox .RejectRole )
24
-
25
71
msg_box .exec ()
26
72
clicked_button = msg_box .clickedButton ()
27
-
28
73
if clicked_button == cancel_button :
29
74
return
30
-
31
75
file_dialog = QFileDialog ()
32
-
76
+ def start_worker (source ):
77
+ progress = QProgressDialog ("Creating symlinks..." , "Cancel" , 0 , 0 )
78
+ progress .setWindowModality (Qt .WindowModal )
79
+ progress .setMinimumDuration (0 )
80
+ worker = SymlinkWorker (source , target_dir )
81
+ main_window = _get_main_window ()
82
+ if main_window and hasattr (main_window , 'databases_tab' ):
83
+ db_tab = main_window .databases_tab
84
+ if hasattr (db_tab , 'docs_model' ) and db_tab .docs_model :
85
+ if hasattr (QFileSystemModel , 'DontWatchForChanges' ):
86
+ db_tab .docs_model .setOption (QFileSystemModel .DontWatchForChanges , True )
87
+ if hasattr (db_tab , 'docs_refresh' ):
88
+ db_tab .docs_refresh .start ()
89
+ progress .canceled .connect (worker .requestInterruption )
90
+ def update_progress (pct ):
91
+ if progress .maximum () == 0 :
92
+ progress .setRange (0 , 100 )
93
+ progress .setValue (pct )
94
+ worker .progress .connect (update_progress )
95
+ def _done (count , errs ):
96
+ if main_window and hasattr (main_window , 'databases_tab' ):
97
+ db_tab = main_window .databases_tab
98
+ if hasattr (db_tab , 'docs_refresh' ):
99
+ db_tab .docs_refresh .stop ()
100
+ if hasattr (db_tab , 'docs_model' ) and db_tab .docs_model :
101
+ if hasattr (db_tab .docs_model , 'refresh' ):
102
+ db_tab .docs_model .refresh ()
103
+ elif hasattr (db_tab .docs_model , 'reindex' ):
104
+ db_tab .docs_model .reindex ()
105
+ if hasattr (QFileSystemModel , 'DontWatchForChanges' ):
106
+ db_tab .docs_model .setOption (QFileSystemModel .DontWatchForChanges , False )
107
+ progress .reset ()
108
+ msg = f"Created { count } symlinks"
109
+ if errs :
110
+ msg += f" – { len (errs )} errors (see console)"
111
+ print (* errs , sep = "\n " )
112
+ QMessageBox .information (None , "Symlinks" , msg )
113
+ worker .finished .connect (_done )
114
+ worker .progress .connect (update_progress )
115
+ worker .start ()
116
+ choose_documents_directory ._symlink_thread = worker
33
117
if clicked_button == dir_button :
34
118
file_dialog .setFileMode (QFileDialog .Directory )
35
119
file_dialog .setOption (QFileDialog .ShowDirsOnly , True )
36
120
selected_dir = file_dialog .getExistingDirectory (None , "Choose Directory for Database" , str (current_dir ))
37
121
if selected_dir :
38
- selected_dir_path = Path (selected_dir )
39
- compatible_files = []
40
- incompatible_files = []
41
-
42
- for file_path in selected_dir_path .iterdir ():
43
- if file_path .is_file ():
44
- if file_path .suffix .lower () in ALLOWED_EXTENSIONS :
45
- compatible_files .append (str (file_path ))
46
- else :
47
- incompatible_files .append (file_path .name )
48
-
49
- if incompatible_files :
50
- if not show_incompatible_files_dialog (incompatible_files ):
51
- return
52
-
53
- if compatible_files :
54
- try :
55
- count , errors = create_symlinks_parallel (compatible_files , target_dir )
56
- if errors :
57
- print ("Errors occurred while creating symlinks:" , errors )
58
- except Exception as e :
59
- print (f"Error creating symlinks: { e } " )
122
+ start_worker (Path (selected_dir ))
60
123
else :
61
124
file_dialog .setFileMode (QFileDialog .ExistingFiles )
62
125
file_paths = file_dialog .getOpenFileNames (None , "Choose Documents and Images for Database" , str (current_dir ))[0 ]
63
126
if file_paths :
64
127
compatible_files = []
65
128
incompatible_files = []
66
-
67
129
for file_path in file_paths :
68
130
path = Path (file_path )
69
131
if path .suffix .lower () in ALLOWED_EXTENSIONS :
70
132
compatible_files .append (str (path ))
71
133
else :
72
134
incompatible_files .append (path .name )
73
-
74
- if incompatible_files :
75
- if not show_incompatible_files_dialog (incompatible_files ):
76
- return
77
-
135
+ if incompatible_files and not show_incompatible_files_dialog (incompatible_files ):
136
+ return
78
137
if compatible_files :
79
- try :
80
- count , errors = create_symlinks_parallel (compatible_files , target_dir )
81
- if errors :
82
- print ("Errors occurred while creating symlinks:" , errors )
83
- except Exception as e :
84
- print (f"Error creating symlinks: { e } " )
85
-
138
+ start_worker (compatible_files )
86
139
def show_incompatible_files_dialog (incompatible_files ):
87
- dialog_text = (
88
- "The following files cannot be added here due to their file extension:\n \n " +
89
- "\n " .join (incompatible_files ) +
90
- "\n \n However, if any of them are audio files you can still add them directly in the Tools Tab."
91
- "\n \n Click 'Ok' to add the compatible documents only (remembering to add audio files separately) or 'Cancel' to back out completely."
92
- )
140
+ dialog_text = ("The following files cannot be added here due to their file extension:\n \n " + "\n " .join (incompatible_files ) + "\n \n However, if any of them are audio files you can still add them directly in the Tools Tab."
141
+ "\n \n Click 'Ok' to add the compatible documents only (remembering to add audio files separately) or 'Cancel' to back out completely." )
93
142
incompatible_dialog = QDialog ()
94
143
incompatible_dialog .resize (800 , 600 )
95
144
incompatible_dialog .setWindowTitle ("Incompatible Files Detected" )
96
-
97
145
layout = QVBoxLayout ()
98
146
text_edit = QTextEdit ()
99
147
text_edit .setReadOnly (True )
@@ -109,18 +157,19 @@ def show_incompatible_files_dialog(incompatible_files):
109
157
ok_button .clicked .connect (incompatible_dialog .accept )
110
158
cancel_button .clicked .connect (incompatible_dialog .reject )
111
159
return incompatible_dialog .exec () == QDialog .Accepted
112
-
113
160
def load_config ():
114
161
with open (CONFIG_FILE , 'r' , encoding = 'utf-8' ) as stream :
115
162
return yaml .safe_load (stream )
116
-
117
163
def select_embedding_model_directory ():
118
164
initial_dir = Path ('Models' ) if Path ('Models' ).exists () else Path .home ()
119
165
chosen_directory = QFileDialog .getExistingDirectory (None , "Select Embedding Model Directory" , str (initial_dir ))
120
-
121
166
if chosen_directory :
122
167
config_file_path = Path (CONFIG_FILE )
123
168
config_data = yaml .safe_load (config_file_path .read_text (encoding = 'utf-8' )) if config_file_path .exists () else {}
124
169
config_data ["EMBEDDING_MODEL_NAME" ] = chosen_directory
125
170
config_file_path .write_text (yaml .dump (config_data ), encoding = 'utf-8' )
126
- print (f"Selected directory: { chosen_directory } " )
171
+ def _get_main_window ():
172
+ for widget in QApplication .topLevelWidgets ():
173
+ if hasattr (widget , 'databases_tab' ):
174
+ return widget
175
+ return None
0 commit comments