Update app.py

app.py

@@ -5,8 +5,7 @@ import os
 import json
 import pandas as pd
 import uuid
-from PIL import Image, ImageDraw # For minimap (using buttons instead now)
-# Corrected import: removed 'sync'
+from PIL import Image, ImageDraw # Not used for minimap anymore, but kept just in case
 from streamlit_js_eval import streamlit_js_eval # For JS communication
 
 # --- Constants ---
@@ -24,7 +23,15 @@ def load_plot_metadata():
     """Scans save dir, sorts plots, calculates metadata."""
     plots = []
     # Ensure consistent sorting, e.g., alphabetically which often aligns with plot_001, plot_002 etc.
-    plot_files = sorted([f for f in os.listdir(SAVE_DIR) if f.endswith(".csv")])
+    try:
+        plot_files = sorted([f for f in os.listdir(SAVE_DIR) if f.endswith(".csv")])
+    except FileNotFoundError:
+        st.error(f"Save directory '{SAVE_DIR}' not found.")
+        return [], 0.0 # Return empty if dir doesn't exist
+    except Exception as e:
+        st.error(f"Error listing save directory '{SAVE_DIR}': {e}")
+        return [], 0.0
+
 
     current_x_offset = 0.0
     for i, filename in enumerate(plot_files):
@@ -54,11 +61,10 @@ def load_plot_objects(filename, x_offset):
              st.warning(f"CSV '{filename}' missing essential columns (type, pos_x/y/z). Skipping.")
              return []
         # Ensure optional columns default to something sensible if missing
-        if 'obj_id' not in df.columns: df['obj_id'] = [str(uuid.uuid4()) for _ in range(len(df))]
-        if 'rot_x' not in df.columns: df['rot_x'] = 0.0
-        if 'rot_y' not in df.columns: df['rot_y'] = 0.0
-        if 'rot_z' not in df.columns: df['rot_z'] = 0.0
-        if 'rot_order' not in df.columns: df['rot_order'] = 'XYZ'
+        # Use vectorized operations for defaults where possible
+        df['obj_id'] = df.get('obj_id', pd.Series([str(uuid.uuid4()) for _ in range(len(df))]))
+        for col, default in [('rot_x', 0.0), ('rot_y', 0.0), ('rot_z', 0.0), ('rot_order', 'XYZ')]:
+             if col not in df.columns: df[col] = default
 
 
         for _, row in df.iterrows():
@@ -171,7 +177,10 @@ with st.sidebar:
         if cols[col_idx].button(button_label, key=f"nav_{plot['filename']}"):
             # Send command to JS to move the player
             target_x = plot['x_offset']
-            streamlit_js_eval(js_code=f"teleportPlayer({target_x});")
+            try:
+                streamlit_js_eval(js_code=f"teleportPlayer({target_x});", key=f"teleport_{plot['filename']}")
+            except Exception as e:
+                 st.error(f"Failed to send teleport command: {e}")
             # No rerun needed here, JS handles the move instantly
 
         col_idx = (col_idx + 1) % max_cols # Cycle through columns
@@ -196,10 +205,13 @@ with st.sidebar:
         key="selected_object_widget" # Use a distinct key for the widget
     )
     # Update session state only if the widget's value actually changes
+    # This change WILL trigger a rerun because Streamlit tracks widget state.
+    # The JS side now handles preserving its state across the resulting reload via sessionStorage.
     if selected_object_type_widget != st.session_state.selected_object:
          st.session_state.selected_object = selected_object_type_widget
-         # Trigger a targeted JS update if needed, or rely on injection next cycle
-         # No rerun needed just for selection if JS reads global state on demand
+         # We don't *need* to force a rerun here, Streamlit handles it.
+         # The important part is that the NEXT run will inject the new selected type,
+         # and the JS will restore placed objects from sessionStorage.
 
 
     st.markdown("---")
@@ -264,7 +276,7 @@ if save_data_from_js is not None: # Process only if data is present
                 st.session_state.new_plot_name = ""
                 # Reset newly placed objects in JS AFTER successful save
                 try:
-                     # This call tells JS to clear its internal 'newlyPlacedObjects' array
+                     # This call tells JS to clear its internal 'newlyPlacedObjects' array AND sessionStorage
                      streamlit_js_eval(js_code="resetNewlyPlacedObjects();", key="reset_js_state")
                 except Exception as js_e:
                      st.warning(f"Could not reset JS state after save: {js_e}")