1+ import json
2+
3+ from PyQt5 .QtGui import QIcon
4+ from PyQt5 .QtWidgets import (
5+ QWidget , QHBoxLayout , QLineEdit , QPushButton , QTreeWidget , QTreeWidgetItem , QStyle
6+ )
7+
8+ from finch .common import s3_session , ObjectType , StringUtils , resource_path
9+
10+
11+ class SearchWidget (QWidget ):
12+ def __init__ (self , main_widget : QWidget ):
13+ super ().__init__ ()
14+ self .main_widget = main_widget
15+ self .icon_type = self ._initialize_icons ()
16+ self ._init_ui ()
17+
18+ def showEvent (self , event ):
19+ """Ensure search input gets focus when the widget is shown."""
20+ super ().showEvent (event )
21+ self .search_input .setFocus ()
22+
23+ def close (self ):
24+ super ().close ()
25+ for idx , action in enumerate (self .main_widget .file_toolbar .actions ()):
26+ if idx in [5 ]:
27+ action .setDisabled (False )
28+ self .main_widget .layout .removeWidget (self )
29+
30+
31+ def _initialize_icons (self ):
32+ """Initialize icon mapping for different object types."""
33+ style = self .style ()
34+ return {
35+ ObjectType .FILE : style .standardIcon (QStyle .SP_FileIcon ),
36+ ObjectType .FOLDER : style .standardIcon (QStyle .SP_DirIcon ),
37+ ObjectType .BUCKET : style .standardIcon (QStyle .SP_DirIcon ),
38+ }
39+
40+ def _init_ui (self ):
41+ """Initialize UI components."""
42+ layout = QHBoxLayout ()
43+ self .search_input = QLineEdit (placeholderText = "Search" )
44+ self .search_input .returnPressed .connect (self ._on_search )
45+ self .search_button = QPushButton ("Search" )
46+ self .search_button .clicked .connect (self ._on_search )
47+ self .close_button = QPushButton ("" )
48+ self .close_button .setIcon (QIcon (resource_path ('img/close.svg' )))
49+ self .close_button .setFlat (True )
50+ self .close_button .setStyleSheet ("QPushButton { background-color: transparent }" )
51+ self .close_button .clicked .connect (self .close )
52+
53+ layout .addWidget (self .search_input )
54+ layout .addWidget (self .search_button )
55+ layout .addWidget (self .close_button )
56+ self .setLayout (layout )
57+
58+ def _on_search (self ):
59+ """Handle search button click."""
60+ search_term = self .search_input .text ()
61+ self .main_widget .tree_widget .clear ()
62+ self ._search_and_populate (search_term )
63+
64+ for i in range (self .main_widget .tree_widget .topLevelItemCount ()):
65+ self ._expand_and_select (self .main_widget .tree_widget .topLevelItem (i ), search_term )
66+
67+ def _search_and_populate (self , search_term ):
68+ """Search S3 and populate the tree widget."""
69+ buckets = self ._get_s3_buckets ()
70+ items = self ._search_s3_objects (buckets , search_term )
71+
72+ for bucket in buckets :
73+ bucket_item = self ._create_tree_item (
74+ name = bucket ['Name' ], object_type = ObjectType .BUCKET , date = bucket ['CreationDate' ]
75+ )
76+
77+ self .main_widget .tree_widget .addTopLevelItem (bucket_item )
78+
79+ bucket_objects = [
80+ (item ['Key' ], item ['Size' ], item ['LastModified' ])
81+ for name , item in items if name == bucket ['Name' ]
82+ ]
83+ tree_structure = self ._build_tree_structure (bucket_objects )
84+ self ._add_items_to_tree (bucket_item , tree_structure )
85+
86+ def _get_s3_buckets (self ):
87+ """Retrieve list of S3 buckets."""
88+ return s3_session .resource .meta .client .list_buckets ()['Buckets' ]
89+
90+ def _search_s3_objects (self , buckets , search_term ):
91+ """Search for objects in S3 matching the search term."""
92+ items = []
93+ for bucket in buckets :
94+ paginator = s3_session .resource .meta .client .get_paginator ('list_objects_v2' )
95+ for obj in paginator .paginate (Bucket = bucket ['Name' ]).search (
96+ f"Contents[?contains(Key, `{ json .dumps (search_term )}
97+ ):
98+ if obj :
99+ items .append ((bucket ['Name' ], obj ))
100+ return items
101+
102+ def _build_tree_structure (self , objects ):
103+ """Build a nested dictionary representing the folder structure."""
104+ tree = {}
105+ for path , size , date in objects :
106+ current = tree
107+ * folders , filename = path .split ('/' )
108+ for folder in folders :
109+ current = current .setdefault (folder , {})
110+ current [filename ] = {"_info" : (size , date )}
111+ return tree
112+
113+ def _add_items_to_tree (self , parent_item , tree_dict ):
114+ """Recursively add items to the tree widget."""
115+ for key , value in tree_dict .items ():
116+ if key == "_info" :
117+ continue
118+
119+ item = self ._create_tree_item (name = key , object_type = ObjectType .FOLDER )
120+ parent_item .addChild (item )
121+
122+ if "_info" in value :
123+ size , date = value ["_info" ]
124+ item = self ._create_tree_item (name = key , object_type = ObjectType .FILE , size = size , date = date )
125+
126+ self ._add_items_to_tree (item , value )
127+
128+ def _expand_and_select (self , item , search_term ):
129+ """Recursively expand and select matching items."""
130+ if item .childCount ():
131+ item .setExpanded (True )
132+ if search_term in item .text (0 ):
133+ item .setSelected (True )
134+ for i in range (item .childCount ()):
135+ self ._expand_and_select (item .child (i ), search_term )
136+
137+ def _create_tree_item (self , name , object_type , size = 0 , date = None ):
138+ """Create a QTreeWidgetItem with the given texts, type, icon, size, and date."""
139+ item = QTreeWidgetItem ()
140+ item .setText (0 , name )
141+ item .setIcon (0 , self .icon_type [object_type ])
142+ item .setText (1 , object_type )
143+ item .setText (2 , StringUtils .format_size (size ))
144+ item .setText (3 , StringUtils .format_datetime (date ))
145+ return item
0 commit comments