Update app.py

app.py

@@ -16,13 +16,13 @@ from reportlab.lib import colors
 class KDChecker:
     def __init__(self):
         self.excel_db = pd.DataFrame()
-        self.cabinet_list = []
         self.known_docs = ["Э3", "В4", "ПЭ3", "ВО", "ТЭ5", "СБ", "С5", "ОЛ", "Э1", "Э4", "Э7", "Д3", "Э6"]
 
     def load_excel_db(self, excel_path):
         print(f"--- Загрузка Excel: {excel_path} ---")
         if excel_path is None:
-            return "Файл не выбран", gr.update(choices=[], value=None)
+            # Возвращаем: Сообщение, Пустой список авторов, Пустой список шкафов
+            return "Файл не выбран", gr.update(choices=[], value=None), gr.update(choices=[], value=None)
 
         all_data = []
         sheets_log = []
@@ -52,6 +52,9 @@ class KDChecker:
                     df = pd.read_excel(excel_path, sheet_name=sheet_name, header=header_row_index)
                     df_subset = df.iloc[:, [cab_col_idx, rem_col_idx]]
                     df_subset.columns = ["Cabinet", "Remark"]
+                    # Добавляем колонку с именем листа (Автором)
+                    df_subset["Author"] = sheet_name 
+                    
                     df_subset["Cabinet"] = df_subset["Cabinet"].ffill()
                     df_subset = df_subset.dropna(subset=["Remark"]).astype(str)
                     df_subset["Cabinet_Clean"] = df_subset["Cabinet"].apply(
@@ -61,25 +64,37 @@ class KDChecker:
                     sheets_log.append(f"'{sheet_name}': {len(df_subset)}")
 
             if not all_data:
-                return "❌ Ошибка: Не найдены заголовки 'Шкаф' и 'Примечание'.", gr.update(choices=[], value=None)
+                return "❌ Ошибка: Не найдены заголовки 'Шкаф' и 'Примечание'.", gr.update(choices=[]), gr.update(choices=[])
 
             self.excel_db = pd.concat(all_data, ignore_index=True)
-            self.cabinet_list = sorted(self.excel_db["Cabinet"].unique().tolist())
+            
+            # Получаем список уникальных авторов (листов)
+            authors_list = sorted(self.excel_db["Author"].unique().tolist())
 
             print(f"Excel загружен. Всего строк: {len(self.excel_db)}")
-            msg = f"✅ База загружена!\nЗаписей: {len(self.excel_db)}\nЛисты: {', '.join(sheets_log)}"
-            return msg, gr.update(choices=self.cabinet_list, value=None, interactive=True)
+            msg = f"✅ База загружена!\nЗаписей: {len(self.excel_db)}\nАвторы: {', '.join(sheets_log)}"
+            
+            # Возвращаем: Статус, Список авторов, Очищаем список шкафов
+            return msg, gr.update(choices=authors_list, value=None, interactive=True), gr.update(choices=[], value=None)
 
         except Exception as e:
             print(f"Ошибка Excel: {e}")
-            return f"❌ Ошибка: {e}", gr.update(choices=[], value=None)
+            return f"❌ Ошибка: {e}", gr.update(choices=[]), gr.update(choices=[])
+
+    def get_cabinets_by_author(self, author_name):
+        """Фильтрует шкафы по выбранному автору"""
+        if self.excel_db.empty or not author_name:
+            return gr.update(choices=[], value=None)
+        
+        # Фильтруем базу по автору
+        filtered_cabs = self.excel_db[self.excel_db["Author"] == author_name]["Cabinet"].unique().tolist()
+        return gr.update(choices=sorted(filtered_cabs), value=None, interactive=True)
 
     def extract_text(self, pdf_path):
         text = ""
         try:
             with pdfplumber.open(pdf_path) as pdf:
-                # Читаем только первые 2 страницы для скорости, обычно номер там
-                # Если нужно читать весь файл - уберите [:2]
+                # Читаем первые 5 страниц для скорости поиска номера
                 for page in pdf.pages[:5]: 
                     text += (page.extract_text() or "") + "\n"
         except Exception as e:
@@ -164,35 +179,27 @@ class KDChecker:
             found_by_method = "manual"
             is_manual = True
         else:
-            # === ОПТИМИЗАЦИЯ: Читаем файлы по очереди и прерываемся, если нашли ===
+            # === ОПТИМИЗАЦИЯ ===
             db_clean_keys = set(self.excel_db["Cabinet_Clean"].tolist())
             
             for file_path in progress.tqdm(files, desc="Поиск номера шкафа"):
                 text = self.extract_text(file_path)
                 pdf_numbers = self.find_all_decimal_numbers(text)
                 
-                # Проверяем, есть ли найденные номера в базе
                 for cand in pdf_numbers:
                     if cand in db_clean_keys:
                         detected_cabinet = cand
                         found_by_method = "number"
-                        break # Прерываем поиск по номерам в текущем файле
+                        break 
                 
                 if found_by_method == "number":
                     print(f"✅ Шкаф найден в файле: {os.path.basename(file_path)}")
-                    break # Прерываем чтение остальных файлов!
-
-            # Если по номерам не нашли, пробуем по имени (тоже не читаем все подряд)
-            if detected_cabinet == "Не определен":
-                # Здесь сложнее, так как имя может быть в любом файле. 
-                # Но если мы уже прочитали часть файлов и не нашли номер, 
-                # скорее всего это ручной ввод.
-                pass 
+                    break 
 
         print(f"Определен шкаф: {detected_cabinet}")
 
         if detected_cabinet == "Не определен":
-            return f"⚠️ Шкаф не опознан автоматически.\nВыберите шкаф вручную.", None
+            return f"⚠️ Шкаф не опознан автоматически.\nВыберите Автора и Шкаф вручную.", None
 
         is_clean_search = (found_by_method == "number")
         remarks = self.get_remarks(detected_cabinet, is_clean_key=is_clean_search)
@@ -200,7 +207,6 @@ class KDChecker:
         if not remarks:
             return f"⚠️ Для шкафа '{detected_cabinet}' нет замечаний в базе.", None
 
-        # Формируем список задач (здесь уже не читаем PDF, только имена файлов)
         processed_count = 0
         for file_path in files:
             fname = os.path.basename(file_path)
@@ -329,14 +335,28 @@ def create_app():
         gr.Markdown("## ✅ Генератор чек-листов КД")
 
         with gr.Row():
+            # --- ЛЕВАЯ КОЛОНКА: БАЗА ---
             with gr.Column(scale=1):
                 gr.Markdown("### 1. База знаний")
                 db_in = gr.File(label="Excel (.xlsx)", type="filepath", elem_classes="compact_file")
+                
+                # Группа ручного выбора
                 with gr.Group():
-                    manual_cab = gr.Dropdown(label="Или шкаф вручную", choices=[], interactive=True)
+                    gr.Markdown("#### Ручной выбор (если автопоиск не сработал):")
+                    # Сначала выбираем автора
+                    author_dd = gr.Dropdown(label="1. Автор (Лист Excel)", choices=[], interactive=True)
+                    # Затем шкаф (зависит от автора)
+                    cabinet_dd = gr.Dropdown(label="2. Выберите шкаф", choices=[], interactive=True)
+                    
                     db_out = gr.Textbox(label="Статус базы", lines=2, max_lines=3, interactive=False)
-                db_in.upload(checker.load_excel_db, inputs=[db_in], outputs=[db_out, manual_cab])
 
+                # Событие загрузки Excel: обновляет статус и список авторов
+                db_in.upload(checker.load_excel_db, inputs=[db_in], outputs=[db_out, author_dd, cabinet_dd])
+                
+                # Событие выбора автора: фильтрует список шкафов
+                author_dd.change(checker.get_cabinets_by_author, inputs=[author_dd], outputs=[cabinet_dd])
+
+            # --- ПРАВАЯ КОЛОНКА: ЧЕРТЕЖИ ---
             with gr.Column(scale=1):
                 gr.Markdown("### 2. Документация")
                 files_in = gr.File(label="Чертежи (PDF)", file_count="multiple", type="filepath", elem_classes="compact_file")
@@ -350,7 +370,7 @@ def create_app():
             with gr.Column(scale=1):
                 res_pdf = gr.File(label="Скачать готовый PDF")
 
-        btn.click(checker.check_files, inputs=[files_in, manual_cab], outputs=[res_txt, res_pdf])
+        btn.click(checker.check_files, inputs=[files_in, cabinet_dd], outputs=[res_txt, res_pdf])
     
     return app