Update app.py

app.py

@@ -169,23 +169,18 @@ def format_play_time(time_str):
         return None
 
 
-# --- UPDATED Helper function to add TMS location column ---
+# --- NEW Helper function to add TMS location column ---
 def add_tms_locations_to_analysis(analysis_df, tms_movie_list):
     locations = []
     for index, row in analysis_df.iterrows():
         movie_title = row['影片']
         found_versions = []
         for tms_movie in tms_movie_list:
-            # FIX 3: Change matching from 'in' to 'startswith'
-            if tms_movie['assert_name'].startswith(movie_title):
+            if movie_title in tms_movie['assert_name']:
+                # Extract version name by removing the base title
                 version_name = tms_movie['assert_name'].replace(movie_title, '').strip()
                 circled_halls = " ".join(sorted([get_circled_number(h) for h in tms_movie['halls']]))
-
-                # FIX 2: Handle empty version name to remove colon
-                if version_name:
-                    found_versions.append(f"{version_name}：{circled_halls}")
-                else:
-                    found_versions.append(circled_halls)
+                found_versions.append(f"{version_name}：{circled_halls}")
 
         locations.append('｜'.join(found_versions))
 
@@ -198,6 +193,7 @@ st.title('影城排片效率与内容分析工具')
 st.write("上传 `影片映出日累计报表.xlsx` 进行效率分析，或点击下方按钮查询 TMS 服务器影片内容。")
 
 uploaded_file = st.file_uploader("请在此处上传 Excel 文件", type=['xlsx', 'xls', 'csv'])
+# NEW: Checkbox for the new feature
 query_tms_for_location = st.checkbox("查询 TMS 找影片所在影厅")
 
 if uploaded_file is not None:
@@ -221,23 +217,12 @@ if uploaded_file is not None:
         prime_time_analysis = process_and_analyze_data(
             df[df['放映时间'].between(pd.to_datetime('14:00:00').time(), pd.to_datetime('21:00:00').time())].copy())
 
+        # --- NEW LOGIC: If checkbox is ticked, fetch data and modify dataframes ---
         if query_tms_for_location:
             with st.spinner("正在关联查询 TMS 服务器..."):
                 _, tms_movie_list = fetch_and_process_server_movies()
                 full_day_analysis = add_tms_locations_to_analysis(full_day_analysis, tms_movie_list)
                 prime_time_analysis = add_tms_locations_to_analysis(prime_time_analysis, tms_movie_list)
-
-                # FIX 1: Reorder columns
-                if '影片所在影厅位置' in full_day_analysis.columns:
-                    cols_full = full_day_analysis.columns.tolist()
-                    cols_full.insert(1, cols_full.pop(cols_full.index('影片所在影厅位置')))
-                    full_day_analysis = full_day_analysis[cols_full]
-
-                if '影片所在影厅位置' in prime_time_analysis.columns:
-                    cols_prime = prime_time_analysis.columns.tolist()
-                    cols_prime.insert(1, cols_prime.pop(cols_prime.index('影片所在影厅位置')))
-                    prime_time_analysis = prime_time_analysis[cols_prime]
-
                 st.toast("TMS 影片位置关联成功！", icon="🔗")
 
         st.markdown("### 全天排片效率分析")