lorax-arg 0.1__py3-none-any.whl

lorax/sockets/file_ops.py

@@ -0,0 +1,200 @@
+"""
+File operation event handlers for Lorax Socket.IO.
+
+Handles load_file, details, and query events.
+"""
+
+import os
+import json
+import asyncio
+from pathlib import Path
+
+from lorax.context import session_manager, BUCKET_NAME, tree_graph_cache, csv_tree_graph_cache
+from lorax.modes import CURRENT_MODE
+from lorax.constants import (
+    UPLOADS_DIR, ERROR_SESSION_NOT_FOUND, ERROR_MISSING_SESSION, ERROR_NO_FILE_LOADED,
+)
+from lorax.cloud.gcs_utils import download_gcs_file
+from lorax.handlers import handle_upload, handle_details
+from lorax.sockets.decorators import require_session
+from lorax.sockets.utils import is_csv_session_file
+
+UPLOAD_DIR = Path(UPLOADS_DIR)
+UPLOAD_DIR.mkdir(exist_ok=True)
+
+
+def register_file_events(sio):
+    """Register file operation socket events."""
+
+    async def background_load_file(sid, data):
+        try:
+            lorax_sid = data.get("lorax_sid")
+            share_sid = data.get("share_sid")
+
+            if not lorax_sid:
+                print(f"⚠️ Missing lorax_sid")
+                await sio.emit("error", {
+                    "code": ERROR_MISSING_SESSION,
+                    "message": "Session ID is missing."
+                }, to=sid)
+                return
+
+            session = await session_manager.get_session(lorax_sid)
+            if not session:
+                print(f"⚠️ Unknown sid {lorax_sid}")
+                await sio.emit("error", {
+                    "code": ERROR_SESSION_NOT_FOUND,
+                    "message": "Session expired. Please refresh the page."
+                }, to=sid)
+                return
+
+            if share_sid and share_sid != lorax_sid:
+                print(f"⚠️ share_sid denied for sid={lorax_sid} target={share_sid}")
+                await sio.emit("error", {
+                    "code": "share_sid_denied",
+                    "message": "Access denied for shared upload."
+                }, to=sid)
+                return
+
+            project = str(data.get("project"))
+            filename = str(data.get("file"))
+
+            # Extract genomic coordinates from client if provided
+            genomiccoordstart = data.get("genomiccoordstart")
+            genomiccoordend = data.get("genomiccoordend")
+            print("lorax_sid", lorax_sid, project, filename)
+            if not filename:
+                print("Missing file param")
+                return
+
+            if project == 'Uploads':
+                target_sid = share_sid if share_sid else lorax_sid
+                if CURRENT_MODE == "local":
+                    file_path = UPLOAD_DIR / project / filename
+                    blob_path = f"{project}/{filename}"
+                else:
+                    file_path = UPLOAD_DIR / project / target_sid / filename
+                    blob_path = f"{project}/{target_sid}/{filename}"
+            else:
+                file_path = UPLOAD_DIR / project / filename
+                blob_path = f"{project}/{filename}"
+
+            if BUCKET_NAME and CURRENT_MODE != "local":
+                if file_path.exists():
+                    print(f"File {file_path} already exists, skipping download.")
+                else:
+                    print(f"Downloading file {file_path} from {BUCKET_NAME}")
+                    await download_gcs_file(BUCKET_NAME, f"{blob_path}", str(file_path))
+            else:
+                print("local mode or no bucket; using local uploads only")
+
+            if not file_path.exists():
+                print("File not found")
+                return
+
+            # Clear TreeGraph cache when loading a new file
+            await tree_graph_cache.clear_session(lorax_sid)
+            await csv_tree_graph_cache.clear_session(lorax_sid)
+
+            session.file_path = str(file_path)
+            await session_manager.save_session(session)
+
+            print("loading file", file_path, os.getpid())
+            ctx = await handle_upload(str(file_path), str(UPLOAD_DIR))
+
+            await sio.emit("status", {
+                "status": "processing-file",
+                "message": "Processing file...",
+                "filename": filename,
+                "project": project
+            }, to=sid)
+
+            # Config is already computed and cached in FileContext
+            config = ctx.config
+
+            if config is None:
+                await sio.emit("error", {"message": "Failed to load file configuration"}, to=sid)
+                return
+
+            # Override initial_position if client provided genomic coordinates
+            if genomiccoordstart is not None and genomiccoordend is not None:
+                try:
+                    config['initial_position'] = [int(genomiccoordstart), int(genomiccoordend)]
+                    print(f"Using client-provided coordinates: [{genomiccoordstart}, {genomiccoordend}]")
+                except (ValueError, TypeError) as e:
+                    print(f"Invalid coordinates, using computed: {e}")
+
+            owner_sid = share_sid if share_sid else lorax_sid
+            await sio.emit("load-file-result", {
+                "message": "File loaded",
+                "sid": sid,
+                "filename": filename,
+                "config": config,
+                "owner_sid": owner_sid
+            }, to=sid)
+
+        except Exception as e:
+            print(f"Load file error: {e}")
+            await sio.emit("error", {"message": str(e)}, to=sid)
+
+    @sio.event
+    async def load_file(sid, data):
+        asyncio.create_task(background_load_file(sid, data))
+
+    @sio.event
+    async def details(sid, data):
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return
+
+            if not session.file_path:
+                print(f"⚠️ No file loaded for session {lorax_sid}")
+                await sio.emit("error", {
+                    "code": ERROR_NO_FILE_LOADED,
+                    "message": "No file loaded. Please load a file first."
+                }, to=sid)
+                return
+
+            if is_csv_session_file(session.file_path):
+                await sio.emit("details-result", {
+                    "data": {"error": "Details are not supported for CSV yet."}
+                }, to=sid)
+                return
+
+            print("fetch details in ", session.sid, os.getpid())
+
+            result = await handle_details(session.file_path, data)
+            await sio.emit("details-result", {"data": json.loads(result)}, to=sid)
+        except Exception as e:
+            print(f"❌ Details error: {e}")
+            await sio.emit("details-result", {"error": str(e)}, to=sid)
+
+    @sio.event
+    async def query(sid, data):
+        """Socket event to query tree nodes."""
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return
+
+            if not session.file_path:
+                print(f"⚠️ No file loaded for session {lorax_sid}")
+                await sio.emit("error", {
+                    "code": ERROR_NO_FILE_LOADED,
+                    "message": "No file loaded. Please load a file first."
+                }, to=sid)
+                return
+
+            value = data.get("value")
+            local_trees = data.get("localTrees", [])
+
+            # Acknowledge the query - the actual tree data is processed by the frontend worker
+            await sio.emit("query-result", {
+                "data": {"value": value, "localTrees": local_trees}
+            }, to=sid)
+        except Exception as e:
+            print(f"❌ Query error: {e}")
+            await sio.emit("query-result", {"error": str(e)}, to=sid)

lorax/sockets/lineage.py

@@ -0,0 +1,307 @@
+"""
+Lineage event handlers for Lorax Socket.IO.
+
+Handles ancestry and descendant operations:
+- get_ancestors_event
+- get_descendants_event
+- search_nodes_by_criteria_event
+- get_subtree_event
+- get_mrca_event
+"""
+
+from lorax.context import tree_graph_cache
+from lorax.handlers import get_or_construct_tree_graph
+from lorax.lineage import (
+    get_ancestors, get_descendants, search_nodes_by_criteria,
+    get_subtree, get_mrca
+)
+from lorax.sockets.decorators import require_session
+from lorax.sockets.utils import is_csv_session_file
+
+
+def register_lineage_events(sio):
+    """Register lineage-related socket events."""
+
+    @sio.event
+    async def get_ancestors_event(sid, data):
+        """Socket event to get ancestors (path to root) for a node.
+
+        Requires the tree to be cached first (call cache_trees or process a layout).
+
+        data: {
+            lorax_sid: str,
+            tree_index: int,
+            node_id: int
+        }
+
+        Returns: {
+            ancestors: [int],           # Node IDs from node to root
+            path: [{node_id, time, x, y}],  # Path coordinates for visualization
+            tree_index: int,
+            query_node: int
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return {"error": "Session not found", "ancestors": [], "path": []}
+
+            if not session.file_path:
+                return {"error": "No file loaded", "ancestors": [], "path": []}
+
+            if is_csv_session_file(session.file_path):
+                return {"error": "Lineage not supported for CSV", "ancestors": [], "path": []}
+
+            tree_index = data.get("tree_index")
+            node_id = data.get("node_id")
+
+            if tree_index is None or node_id is None:
+                return {"error": "Missing tree_index or node_id", "ancestors": [], "path": []}
+
+            # Ensure tree is cached
+            await get_or_construct_tree_graph(
+                session.file_path,
+                int(tree_index),
+                lorax_sid,
+                tree_graph_cache
+            )
+
+            result = await get_ancestors(
+                tree_graph_cache,
+                lorax_sid,
+                int(tree_index),
+                int(node_id)
+            )
+            return result
+        except Exception as e:
+            print(f"❌ Get ancestors error: {e}")
+            return {"error": str(e), "ancestors": [], "path": []}
+
+    @sio.event
+    async def get_descendants_event(sid, data):
+        """Socket event to get all descendants of a node.
+
+        data: {
+            lorax_sid: str,
+            tree_index: int,
+            node_id: int,
+            include_tips_only: bool  # Optional, default False
+        }
+
+        Returns: {
+            descendants: [int],
+            tips: [int],
+            total_descendants: int,
+            tree_index: int,
+            query_node: int
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return {"error": "Session not found", "descendants": [], "tips": []}
+
+            if not session.file_path:
+                return {"error": "No file loaded", "descendants": [], "tips": []}
+
+            if is_csv_session_file(session.file_path):
+                return {"error": "Lineage not supported for CSV", "descendants": [], "tips": []}
+
+            tree_index = data.get("tree_index")
+            node_id = data.get("node_id")
+            include_tips_only = data.get("include_tips_only", False)
+
+            if tree_index is None or node_id is None:
+                return {"error": "Missing tree_index or node_id", "descendants": [], "tips": []}
+
+            # Ensure tree is cached
+            await get_or_construct_tree_graph(
+                session.file_path,
+                int(tree_index),
+                lorax_sid,
+                tree_graph_cache
+            )
+
+            result = await get_descendants(
+                tree_graph_cache,
+                lorax_sid,
+                int(tree_index),
+                int(node_id),
+                include_tips_only=include_tips_only
+            )
+            return result
+        except Exception as e:
+            print(f"❌ Get descendants error: {e}")
+            return {"error": str(e), "descendants": [], "tips": []}
+
+    @sio.event
+    async def search_nodes_by_criteria_event(sid, data):
+        """Socket event to search nodes by criteria (time, tip status, etc).
+
+        data: {
+            lorax_sid: str,
+            tree_index: int,
+            criteria: {
+                min_time: float,      # Optional
+                max_time: float,      # Optional
+                is_tip: bool,         # Optional: True for tips, False for internal
+                has_children: bool,   # Optional: inverse of is_tip
+                node_ids: [int]       # Optional: filter to these nodes
+            }
+        }
+
+        Returns: {
+            matches: [int],
+            positions: [{node_id, x, y, time}],
+            total_matches: int
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return {"error": "Session not found", "matches": [], "positions": []}
+
+            if not session.file_path:
+                return {"error": "No file loaded", "matches": [], "positions": []}
+
+            if is_csv_session_file(session.file_path):
+                return {"error": "Search not supported for CSV", "matches": [], "positions": []}
+
+            tree_index = data.get("tree_index")
+            criteria = data.get("criteria", {})
+
+            if tree_index is None:
+                return {"error": "Missing tree_index", "matches": [], "positions": []}
+
+            # Ensure tree is cached
+            await get_or_construct_tree_graph(
+                session.file_path,
+                int(tree_index),
+                lorax_sid,
+                tree_graph_cache
+            )
+
+            result = await search_nodes_by_criteria(
+                tree_graph_cache,
+                lorax_sid,
+                int(tree_index),
+                criteria
+            )
+            return result
+        except Exception as e:
+            print(f"❌ Search nodes error: {e}")
+            return {"error": str(e), "matches": [], "positions": []}
+
+    @sio.event
+    async def get_subtree_event(sid, data):
+        """Socket event to get the complete subtree rooted at a node.
+
+        data: {
+            lorax_sid: str,
+            tree_index: int,
+            root_node_id: int
+        }
+
+        Returns: {
+            nodes: [{node_id, parent_id, x, y, time, is_tip}],
+            edges: [{parent, child}],
+            total_nodes: int
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return {"error": "Session not found", "nodes": [], "edges": []}
+
+            if not session.file_path:
+                return {"error": "No file loaded", "nodes": [], "edges": []}
+
+            if is_csv_session_file(session.file_path):
+                return {"error": "Subtree not supported for CSV", "nodes": [], "edges": []}
+
+            tree_index = data.get("tree_index")
+            root_node_id = data.get("root_node_id")
+
+            if tree_index is None or root_node_id is None:
+                return {"error": "Missing tree_index or root_node_id", "nodes": [], "edges": []}
+
+            # Ensure tree is cached
+            await get_or_construct_tree_graph(
+                session.file_path,
+                int(tree_index),
+                lorax_sid,
+                tree_graph_cache
+            )
+
+            result = await get_subtree(
+                tree_graph_cache,
+                lorax_sid,
+                int(tree_index),
+                int(root_node_id)
+            )
+            return result
+        except Exception as e:
+            print(f"❌ Get subtree error: {e}")
+            return {"error": str(e), "nodes": [], "edges": []}
+
+    @sio.event
+    async def get_mrca_event(sid, data):
+        """Socket event to find the Most Recent Common Ancestor of multiple nodes.
+
+        data: {
+            lorax_sid: str,
+            tree_index: int,
+            node_ids: [int]  # At least 2 nodes
+        }
+
+        Returns: {
+            mrca: int,              # Node ID of MRCA
+            mrca_time: float,
+            mrca_position: {x, y},
+            tree_index: int,
+            query_nodes: [int]
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return {"error": "Session not found", "mrca": None}
+
+            if not session.file_path:
+                return {"error": "No file loaded", "mrca": None}
+
+            if is_csv_session_file(session.file_path):
+                return {"error": "MRCA not supported for CSV", "mrca": None}
+
+            tree_index = data.get("tree_index")
+            node_ids = data.get("node_ids", [])
+
+            if tree_index is None:
+                return {"error": "Missing tree_index", "mrca": None}
+
+            if len(node_ids) < 2:
+                return {"error": "Need at least 2 nodes", "mrca": None}
+
+            # Ensure tree is cached
+            await get_or_construct_tree_graph(
+                session.file_path,
+                int(tree_index),
+                lorax_sid,
+                tree_graph_cache
+            )
+
+            result = await get_mrca(
+                tree_graph_cache,
+                lorax_sid,
+                int(tree_index),
+                [int(n) for n in node_ids]
+            )
+            return result
+        except Exception as e:
+            print(f"❌ Get MRCA error: {e}")
+            return {"error": str(e), "mrca": None}