Update app.py

app.py

@@ -5,62 +5,77 @@ import os
 import re
 import base64
 
-# 1️⃣ 저장소 설정
+# --------------------------
+# 1️⃣ 기본 설정
+# --------------------------
 UPLOAD_DIR = "uploaded_dbs"
-if not os.path.exists(UPLOAD_DIR):
-    os.makedirs(UPLOAD_DIR)
+os.makedirs(UPLOAD_DIR, exist_ok=True)
 
 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_string = base64.b64encode(blob_data).decode('utf-8')
-            return f'<img src="data:image/png;base64,{encoded_string}" width="200" />'
+            encoded = base64.b64encode(blob_data).decode('utf-8')
+            return f'<img src="data:image/png;base64,{encoded}" width="200"/>'
         return str(blob_data)
     except:
         return str(blob_data)
 
 def natural_sort_key(s):
-    if s is None: return []
-    return [int(text) if text.isdigit() else text.lower()
-            for text in re.split(r'(\d+)', str(s))]
+    return [int(t) if t.isdigit() else t.lower()
+            for t in re.split(r'(\d+)', str(s))] if s else []
 
 def refresh_registry_data():
     global db_registry
     db_registry = []
-    db_files = [f for f in os.listdir(UPLOAD_DIR) if f.endswith(".db")]
 
-    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]
-            })
+    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()
 
-    return sorted(list(set([db["standard"] for db in db_registry])))
+    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
 
 # --------------------------
-# 드롭다운 업데이트
+# 3️⃣ 드롭다운
 # --------------------------
 def on_load():
-    standards = refresh_registry_data()
-    return gr.Dropdown(choices=standards)
+    return gr.Dropdown(choices=refresh_registry_data())
 
 def update_version_dd(standard):
     if not standard:
         return gr.Dropdown(choices=[])
-    versions = sorted(list(set([
-        db["version"] for db in db_registry if db["standard"] == standard
-    ])))
+    versions = sorted({db["version"] for db in db_registry if db["standard"] == standard})
     return gr.Dropdown(choices=versions)
 
 def update_category_dd(standard, version):
@@ -70,29 +85,18 @@ def update_category_dd(standard, version):
     choices = ["ALL"]
 
     try:
-        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)
+        conn = sqlite3.connect(get_db_path(standard, version))
+        main_t, col_map = get_main_table_and_columns(conn, standard, version)
 
+        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():
-                c = str(row[ch])
-                cat = str(row[ca])
-                choices.append(f"{c}.{cat}")
+                ch = str(row.iloc[0])
+                ca = str(row.iloc[1])
+                choices.append(f"{ch}.{ca}")
 
         conn.close()
     except:
@@ -101,77 +105,45 @@ 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:
-        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
-
-        # --------------------------
-        # 조회 로직
-        # --------------------------
+        conn = sqlite3.connect(get_db_path(standard, version))
+        main_t, col_map = get_main_table_and_columns(conn, standard, version)
+
         if selection == "ALL":
             df = pd.read_sql(f"SELECT * FROM [{main_t}]", conn)
 
-        elif "." in selection and real_ch and real_cat:
+        elif "." in selection and 'chapter' in col_map and 'category' in col_map:
             ch, ca = selection.split(".", 1)
             df = pd.read_sql(
-                f"SELECT * FROM [{main_t}] WHERE CAST([{real_ch}] AS TEXT)=? AND CAST([{real_cat}] AS TEXT)=?",
+                f"""
+                SELECT * FROM [{main_t}]
+                WHERE CAST([{col_map['chapter']}] AS TEXT)=?
+                AND CAST([{col_map['category']}] AS TEXT)=?
+                """,
                 conn,
                 params=[ch, ca]
             )
-
         else:
             df = pd.read_sql(f"SELECT * FROM [{main_t}]", conn)
 
         conn.close()
 
-        # 🔥 데이터 없을 때
-        if df is None or df.empty:
+        if df.empty:
             return pd.DataFrame({"Info": ["데이터 없음"]})
 
-        # 🔥 컬럼 정리
+        # 공통 처리
         df.columns = [c.lower().strip() for c in df.columns]
 
-        # 🔥 필요한 컬럼만
-        if 'section' in df.columns and 'description' in df.columns:
+        if {'section', 'description'}.issubset(df.columns):
             df = df[['section', 'description']]
 
-        # 🔥 이미지 처리
-        for col in df.columns:
-            df[col] = df[col].apply(blob_to_base64_html)
+        df = df.applymap(blob_to_base64_html)
 
         return df
 
@@ -179,7 +151,7 @@ def display_data(standard, version, selection):
         return pd.DataFrame({"Error": [str(e)]})
 
 # --------------------------
-# 🔥 통합 조회 (핵심)
+# 5️⃣ 통합 조회
 # --------------------------
 def unified_search(bs, bv, bc, cs, cv, cc):
 
@@ -188,38 +160,31 @@ def unified_search(bs, bv, bc, cs, cv, cc):
         return display_data(bs, bv, bc)
 
     # 비교
-    if bs and bv and bc and cs and cv and cc:
+    if all([bs, bv, bc, cs, cv, cc]):
         df_base = display_data(bs, bv, bc)
         df_comp = display_data(cs, cv, cc)
 
-        if df_base is None or df_comp is None:
+        if df_base.empty or df_comp.empty:
             return pd.DataFrame({"Error": ["데이터 없음"]})
 
-        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")
-            )
+        df_base.columns = [c.lower() for c in df_base.columns]
+        df_comp.columns = [c.lower() for c in df_comp.columns]
 
-            return merged
+        if 'section' not in df_base or 'section' not in df_comp:
+            return pd.DataFrame({"Error": ["section 컬럼 없음"]})
 
-        except Exception as e:
-            return pd.DataFrame({"Error": [str(e)]})
+        return pd.merge(
+            df_base,
+            df_comp,
+            on="section",
+            how="outer",
+            suffixes=("_base", "_comp")
+        )
 
     return pd.DataFrame({"Info": ["선택을 완료하세요"]})
 
 # --------------------------
-# UI 구성
+# 6️⃣ UI
 # --------------------------
 with gr.Blocks() as demo:
 
@@ -227,9 +192,7 @@ with gr.Blocks() as demo:
 
     with gr.Row():
 
-        # 왼쪽
         with gr.Column(scale=1):
-
             with gr.Group():
                 gr.Markdown("### 📌 기준 법규")
                 base_standard = gr.Dropdown(label="Standard")
@@ -242,21 +205,18 @@ 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)
 
@@ -274,7 +234,7 @@ with gr.Blocks() as demo:
     )
 
 # --------------------------
-# 실행
+# 7️⃣ 실행
 # --------------------------
 if __name__ == "__main__":
     demo.launch(