Update app.py

app.py

@@ -5,77 +5,62 @@ import os
 import re
 import base64
 
-# --------------------------
-# 1️⃣ 기본 설정
-# --------------------------
+# 1️⃣ 저장소 설정
 UPLOAD_DIR = "uploaded_dbs"
-os.makedirs(UPLOAD_DIR, exist_ok=True)
+if not os.path.exists(UPLOAD_DIR):
+    os.makedirs(UPLOAD_DIR)
 
 db_registry = []
 
 # --------------------------
-# 2️⃣ 유틸 함수
+# 유틸 함수 (그대로 유지)
 # --------------------------
 def blob_to_base64_html(blob_data):
     if blob_data is None or pd.isna(blob_data):
         return ""
     try:
         if isinstance(blob_data, (bytes, bytearray)):
-            encoded = base64.b64encode(blob_data).decode('utf-8')
-            return f'<img src="data:image/png;base64,{encoded}" width="200"/>'
+            encoded_string = base64.b64encode(blob_data).decode('utf-8')
+            return f'<img src="data:image/png;base64,{encoded_string}" width="200" />'
         return str(blob_data)
     except:
         return str(blob_data)
 
 def natural_sort_key(s):
-    return [int(t) if t.isdigit() else t.lower()
-            for t in re.split(r'(\d+)', str(s))] if s else []
+    if s is None: return []
+    return [int(text) if text.isdigit() else text.lower()
+            for text in re.split(r'(\d+)', str(s))]
 
 def refresh_registry_data():
     global db_registry
     db_registry = []
+    db_files = [f for f in os.listdir(UPLOAD_DIR) if f.endswith(".db")]
 
-    for f in os.listdir(UPLOAD_DIR):
-        if f.endswith(".db"):
-            name = f.replace(".db", "")
-            parts = name.split("_")
-            if len(parts) >= 2:
-                db_registry.append({
-                    "path": os.path.join(UPLOAD_DIR, f),
-                    "standard": parts[0],
-                    "version": parts[1]
-                })
-
-    return sorted({db["standard"] for db in db_registry})
-
-def get_db_path(standard, version):
-    return next(db["path"] for db in db_registry
-                if db["standard"] == standard and db["version"] == version)
-
-def get_main_table_and_columns(conn, standard, version):
-    tables = pd.read_sql(
-        "SELECT name FROM sqlite_master WHERE type='table';", conn
-    )['name'].tolist()
+    for filename in db_files:
+        name_only = filename.replace(".db", "")
+        parts = name_only.split("_")
+        if len(parts) >= 2:
+            db_registry.append({
+                "path": os.path.join(UPLOAD_DIR, filename),
+                "standard": parts[0],
+                "version": parts[1]
+            })
 
-    main_t = f"{standard}_{version}"
-    if main_t not in tables:
-        main_t = tables[0]
-
-    cols = pd.read_sql(f"PRAGMA table_info([{main_t}])", conn)['name'].tolist()
-    col_map = {c.lower().strip(): c for c in cols}
-
-    return main_t, col_map
+    return sorted(list(set([db["standard"] for db in db_registry])))
 
 # --------------------------
-# 3️⃣ 드롭다운
+# 드롭다운 업데이트
 # --------------------------
 def on_load():
-    return gr.Dropdown(choices=refresh_registry_data())
+    standards = refresh_registry_data()
+    return gr.Dropdown(choices=standards)
 
 def update_version_dd(standard):
     if not standard:
         return gr.Dropdown(choices=[])
-    versions = sorted({db["version"] for db in db_registry if db["standard"] == standard})
+    versions = sorted(list(set([
+        db["version"] for db in db_registry if db["standard"] == standard
+    ])))
     return gr.Dropdown(choices=versions)
 
 def update_category_dd(standard, version):
@@ -85,18 +70,29 @@ def update_category_dd(standard, version):
     choices = ["ALL"]
 
     try:
-        conn = sqlite3.connect(get_db_path(standard, version))
-        main_t, col_map = get_main_table_and_columns(conn, standard, version)
+        target_db = next(db for db in db_registry if db["standard"] == standard and db["version"] == version)
+        conn = sqlite3.connect(target_db["path"])
+
+        tables = pd.read_sql("SELECT name FROM sqlite_master WHERE type='table';", conn)['name'].tolist()
+        main_t = f"{standard}_{version}"
+
+        if main_t not in tables:
+            main_t = tables[0]
+
+        cols_info = pd.read_sql(f"PRAGMA table_info([{main_t}])", conn)
+        cols = cols_info['name'].tolist()
+        lower_cols = [c.lower() for c in cols]
+
+        if 'chapter' in lower_cols and 'category' in lower_cols:
+            ch = cols[lower_cols.index('chapter')]
+            ca = cols[lower_cols.index('category')]
+
+            df = pd.read_sql(f"SELECT DISTINCT [{ch}], [{ca}] FROM [{main_t}]", conn)
 
-        if 'chapter' in col_map and 'category' in col_map:
-            df = pd.read_sql(
-                f"SELECT DISTINCT [{col_map['chapter']}], [{col_map['category']}] FROM [{main_t}]",
-                conn
-            )
             for _, row in df.iterrows():
-                ch = str(row.iloc[0])
-                ca = str(row.iloc[1])
-                choices.append(f"{ch}.{ca}")
+                c = str(row[ch])
+                cat = str(row[ca])
+                choices.append(f"{c}.{cat}")
 
         conn.close()
     except:
@@ -105,45 +101,77 @@ def update_category_dd(standard, version):
     return gr.Dropdown(choices=choices)
 
 # --------------------------
-# 4️⃣ 데이터 조회
+# 데이터 조회
 # --------------------------
 def display_data(standard, version, selection):
     if not all([standard, version, selection]):
         return pd.DataFrame({"Info": ["선택을 완료하세요"]})
 
     try:
-        conn = sqlite3.connect(get_db_path(standard, version))
-        main_t, col_map = get_main_table_and_columns(conn, standard, version)
-
+        target_db = next(db for db in db_registry if db["standard"] == standard and db["version"] == version)
+        conn = sqlite3.connect(target_db["path"])
+
+        # 🔥 테이블 찾기 (안전하게)
+        tables = pd.read_sql("SELECT name FROM sqlite_master WHERE type='table';", conn)['name'].tolist()
+        main_t = f"{standard}_{version}"
+
+        if main_t not in tables:
+            main_t = tables[0]
+
+        # 🔥 컬럼 정보 가져오기
+        cols_info = pd.read_sql(f"PRAGMA table_info([{main_t}])", conn)
+        cols = cols_info['name'].tolist()
+        lower_cols = [c.lower().strip() for c in cols]
+
+        # 🔥 실제 컬럼명 찾기
+        real_section = None
+        real_desc = None
+        real_ch = None
+        real_cat = None
+
+        for c, lc in zip(cols, lower_cols):
+            if lc == "section":
+                real_section = c
+            elif lc == "description":
+                real_desc = c
+            elif lc == "chapter":
+                real_ch = c
+            elif lc == "category":
+                real_cat = c
+
+        # --------------------------
+        # 조회 로직
+        # --------------------------
         if selection == "ALL":
             df = pd.read_sql(f"SELECT * FROM [{main_t}]", conn)
 
-        elif "." in selection and 'chapter' in col_map and 'category' in col_map:
+        elif "." in selection and real_ch and real_cat:
             ch, ca = selection.split(".", 1)
             df = pd.read_sql(
-                f"""
-                SELECT * FROM [{main_t}]
-                WHERE CAST([{col_map['chapter']}] AS TEXT)=?
-                AND CAST([{col_map['category']}] AS TEXT)=?
-                """,
+                f"SELECT * FROM [{main_t}] WHERE CAST([{real_ch}] AS TEXT)=? AND CAST([{real_cat}] AS TEXT)=?",
                 conn,
                 params=[ch, ca]
             )
+
         else:
             df = pd.read_sql(f"SELECT * FROM [{main_t}]", conn)
 
         conn.close()
 
-        if df.empty:
+        # 🔥 데이터 없을 때
+        if df is None or df.empty:
             return pd.DataFrame({"Info": ["데이터 없음"]})
 
-        # 공통 처리
+        # 🔥 컬럼 정리
         df.columns = [c.lower().strip() for c in df.columns]
 
-        if {'section', 'description'}.issubset(df.columns):
+        # 🔥 필요한 컬럼만
+        if 'section' in df.columns and 'description' in df.columns:
             df = df[['section', 'description']]
 
-        df = df.applymap(blob_to_base64_html)
+        # 🔥 이미지 처리
+        for col in df.columns:
+            df[col] = df[col].apply(blob_to_base64_html)
 
         return df
 
@@ -151,7 +179,7 @@ def display_data(standard, version, selection):
         return pd.DataFrame({"Error": [str(e)]})
 
 # --------------------------
-# 5️⃣ 통합 조회
+# 🔥 통합 조회 (핵심)
 # --------------------------
 def unified_search(bs, bv, bc, cs, cv, cc):
 
@@ -160,31 +188,38 @@ def unified_search(bs, bv, bc, cs, cv, cc):
         return display_data(bs, bv, bc)
 
     # 비교
-    if all([bs, bv, bc, cs, cv, cc]):
+    if bs and bv and bc and cs and cv and cc:
         df_base = display_data(bs, bv, bc)
         df_comp = display_data(cs, cv, cc)
 
-        if df_base.empty or df_comp.empty:
+        if df_base is None or df_comp is None:
             return pd.DataFrame({"Error": ["데이터 없음"]})
 
-        df_base.columns = [c.lower() for c in df_base.columns]
-        df_comp.columns = [c.lower() for c in df_comp.columns]
+        if 'section' not in df_base.columns or 'section' not in df_comp.columns:
+            return pd.DataFrame({
+                "Error": ["section 컬럼이 없습니다"],
+                "base_columns": [list(df_base.columns)],
+                "comp_columns": [list(df_comp.columns)]
+            })
+
+        try:
+            merged = pd.merge(
+                df_base,
+                df_comp,
+                on="section",
+                how="outer",
+                suffixes=("_base", "_comp")
+            )
 
-        if 'section' not in df_base or 'section' not in df_comp:
-            return pd.DataFrame({"Error": ["section 컬럼 없음"]})
+            return merged
 
-        return pd.merge(
-            df_base,
-            df_comp,
-            on="section",
-            how="outer",
-            suffixes=("_base", "_comp")
-        )
+        except Exception as e:
+            return pd.DataFrame({"Error": [str(e)]})
 
     return pd.DataFrame({"Info": ["선택을 완료하세요"]})
 
 # --------------------------
-# 6️⃣ UI
+# UI 구성
 # --------------------------
 with gr.Blocks() as demo:
 
@@ -192,7 +227,9 @@ with gr.Blocks() as demo:
 
     with gr.Row():
 
+        # 왼쪽
         with gr.Column(scale=1):
+
             with gr.Group():
                 gr.Markdown("### 📌 기준 법규")
                 base_standard = gr.Dropdown(label="Standard")
@@ -205,18 +242,21 @@ with gr.Blocks() as demo:
                 comp_version = gr.Dropdown(label="Version")
                 comp_category = gr.Dropdown(label="Category")
 
+        # 오른쪽
         with gr.Column(scale=4):
             gr.Markdown("### 📊 내용")
+
             output_df = gr.Dataframe(
                 wrap=True,
                 interactive=False,
                 datatype="html"
             )
 
+
     with gr.Row():
         search_btn = gr.Button("🔍 조회")
 
-    # 이벤트
+    # 이벤트 연결
     demo.load(on_load, None, base_standard)
     demo.load(on_load, None, comp_standard)
 
@@ -234,7 +274,7 @@ with gr.Blocks() as demo:
     )
 
 # --------------------------
-# 7️⃣ 실행
+# 실행
 # --------------------------
 if __name__ == "__main__":
     demo.launch(