ops.py

import bpy
from bpy.props import StringProperty
import time
from bpy_extras.io_utils import ImportHelper
from .csv import load_csv
from .parsers import update_obj_from_csv
from .exporters import from_blender_to_polars_df, from_polars_df_to_csv
from pathlib import Path
import csv
import os


# based on the blender docs: https://docs.blender.org/api/current/bpy.types.FileHandler.html#basic-filehandler-for-operator-that-imports-just-one-file
# and tweaked with this prompt: https://chatgpt.com/share/675b0831-354c-8013-bae0-9bb91d527f32


# Operator for the button and drag-and-drop
class ImportCsvPolarsOperator(bpy.types.Operator, ImportHelper):
    bl_idname = "import_scene.import_csv_polars"
    bl_label = "Import CSV (Polars)"
    bl_options = {"PRESET", "UNDO"}

    # ImportHelper mix-in provides 'filepath' by default, but we redefine it here
    # to use SKIP_SAVE, allowing drag-and-drop to work properly.
    filepath: StringProperty(subtype="FILE_PATH", options={"SKIP_SAVE"})  # type: ignore

    filename_ext = ".csv"
    filter_glob: StringProperty(  # type: ignore
        default="*.csv",
        options={"HIDDEN"},
        maxlen=255,
    )

    def execute(self, context):
        # Ensure the filepath is a CSV file

        if not self.filepath.lower().endswith(".csv"):
            self.report({"WARNING"}, "Selected file is not a CSV")
            return {"CANCELLED"}

        start_time = time.perf_counter()

        bob = load_csv(filepath=self.filepath)
        bob.csv.filepath = self.filepath

        elapsed_time_ms = (time.perf_counter() - start_time) * 1000

        self.report(
            {"INFO"},
            f" 🐻‍❄️ 📥  Added {bob.name} in {elapsed_time_ms:.2f} ms",
        )
        return {"FINISHED"}

    def invoke(self, context, event):
        # If the filepath is already set (e.g. drag-and-drop), execute directly
        if self.filepath:
            return self.execute(context)
        # Otherwise, show the file browser
        context.window_manager.fileselect_add(self)
        return {"RUNNING_MODAL"}


# File Handler for drag-and-drop support
class CSV_FH_import(bpy.types.FileHandler):
    bl_idname = "CSV_FH_import"
    bl_label = "File handler for CSV import"
    bl_import_operator = "import_scene.import_csv_polars"
    bl_file_extensions = ".csv"

    @classmethod
    def poll_drop(cls, context):
        # Allow drag-and-drop
        return context.area


class CSV_OP_ReloadData(bpy.types.Operator):
    bl_idname = "csv.reload_data"
    bl_label = "Reload Data"
    bl_options = {"REGISTER", "UNDO"}
    bl_description = (
        "Reload the imported data file with updated values into this object"
    )

    filepath: StringProperty(  # type: ignore
        subtype="FILE_PATH", name="File Path", description="Path to the CSV file"
    )

    def execute(self, context):
        obj: bpy.types.Object = bpy.context.active_object  # type: ignore
        n_points = len(obj.data.vertices)
        update_obj_from_csv(obj, self.filepath)
        message = f"Reloaded data for {len(obj.data.vertices)} points"
        if len(obj.data.vertices) != n_points:
            message += f" (was {n_points})"
        self.report({"INFO"}, message=message)
        return {"FINISHED"}


def hot_reload_timer():
    # The actual hot reload logic
    for obj in bpy.data.objects:
        path = Path(obj.csv.filepath)
        if not obj.csv.hot_reload:
            continue
        if not path.exists():
            obj.csv.hot_reload = False
            continue
        if obj.csv.last_loaded_time < path.stat().st_mtime:
            update_obj_from_csv(obj, str(path))
            obj.csv.last_loaded_time = int(time.time())

    return 1.0  # Run again in 1 second


class CSV_OT_ToggleHotReload(bpy.types.Operator):
    bl_idname = "csv.toggle_hot_reload"
    bl_label = "Toggle Hot Reload"
    bl_description = (
        "Enable or disable hot reloading of the data from the imported file"
    )

    def execute(self, context):
        obj = context.active_object
        if obj.csv.hot_reload:
            try:
                bpy.app.timers.unregister(hot_reload_timer)
            except Exception as e:
                print(e)
            obj.csv.hot_reload = False
            self.report({"INFO"}, "Hot reload stopped")
        else:
            bpy.app.timers.register(hot_reload_timer)
            obj.csv.hot_reload = True
            self.report({"INFO"}, "Hot reload started")
        return {"FINISHED"}

class CSV_OT_ExportData(bpy.types.Operator):
    bl_idname = "csv.export_data"
    bl_label = "Export Data"
    bl_options = {"REGISTER", "UNDO"}
    bl_description = "Export mesh attribute data to file"

    export_type: bpy.props.EnumProperty(  # type: ignore
        name="Export Type",
        description="Type of file to export",
        items=[
            ('CSV', "CSV", "Export as CSV"),
            ('JSON', "JSON", "Export as JSON"), 
            ('PARQUET', "Parquet", "Export as Parquet")
        ],
        default='CSV'
    )

    def execute(self, context):
        scene = context.scene
        export_path = bpy.path.abspath(scene.csv_export.export_path)
        export_object = scene.csv_export.export_object
        
        # Check if an object is selected
        if export_object is None:
            self.report({"WARNING"}, "No object selected for export")
            return {"CANCELLED"}
        
        # Print the object name
        print(f"Exporting data from object: {export_object.name}")
        
        # Create directory if it doesn't exist
        directory = os.path.dirname(export_path)
        if directory and not os.path.exists(directory):
            try:
                os.makedirs(directory)
            except OSError as e:
                self.report({"ERROR"}, f"Failed to create directory: {e}")
                return {"CANCELLED"}
        
        try:
            # Convert Blender object to Polars DataFrame
            df = from_blender_to_polars_df(export_object)
            
            # Export DataFrame based on selected type
            if self.export_type == 'CSV':
                from_polars_df_to_csv(df, export_path)
            elif self.export_type == 'JSON':
                df.write_json(export_path)
            else:  # PARQUET
                df.write_parquet(export_path)
            
            self.report({"INFO"}, f"{self.export_type} file exported to: {export_path} for object: {export_object.name} ({len(df)} rows, {len(df.columns)} columns)")
            return {"FINISHED"}
            
        except Exception as e:
            self.report({"ERROR"}, f"Failed to export {self.export_type} file: {e}")
            return {"CANCELLED"}


CLASSES = (
    ImportCsvPolarsOperator,
    CSV_FH_import,
    CSV_OP_ReloadData,
    CSV_OT_ToggleHotReload,
    CSV_OT_ExportData,
)