Backend abstraction: introduce DataConnection/ProcessConnection ABCs

disruption_py/core/physics_method/params.py

@@ -11,7 +11,7 @@
 from loguru import logger
 
 from disruption_py.core.utils.misc import shot_msg_patch, to_tuple
-from disruption_py.inout.mds import MDSConnection
+from disruption_py.inout.base import DataConnection
 from disruption_py.machine.tokamak import Tokamak
 
 
@@ -25,7 +25,7 @@ class PhysicsMethodParams:
     shot_id: int
     tokamak: Tokamak
     disruption_time: float
-    mds_conn: MDSConnection
+    mds_conn: DataConnection
     times: np.ndarray
 
     def __post_init__(self):

disruption_py/core/retrieval_manager.py

@@ -11,7 +11,8 @@
 from disruption_py.core.physics_method.params import PhysicsMethodParams
 from disruption_py.core.physics_method.runner import populate_shot
 from disruption_py.core.utils.misc import shot_msg
-from disruption_py.inout.mds import MDSConnection, ProcessMDSConnection, mdsExceptions
+from disruption_py.inout.base import DataConnection, ProcessConnection
+from disruption_py.inout.mds import MDSConnection, mdsExceptions
 from disruption_py.inout.sql import ShotDatabase
 from disruption_py.machine.tokamak import Tokamak
 from disruption_py.settings.nickname_setting import NicknameSettingParams
@@ -29,15 +30,15 @@ class RetrievalManager:
         The tokamak instance.
     process_database : ShotDatabase
         The SQL database
-    process_mds_conn : ProcessMDSConnection
-        The MDS connection
+    process_mds_conn : ProcessConnection
+        The process-level data connection
     """
 
     def __init__(
         self,
         tokamak: Tokamak,
         process_database: ShotDatabase,
-        process_mds_conn: ProcessMDSConnection,
+        process_mds_conn: ProcessConnection,
     ):
         """
         Parameters
@@ -46,8 +47,8 @@ def __init__(
             The tokamak instance.
         process_database : ShotDatabase
             The SQL database.
-        process_mds_conn : ProcessMDSConnection
-            The MDS connection.
+        process_mds_conn : ProcessConnection
+            The process-level data connection.
         """
         self.tokamak = tokamak
         self.process_database = process_database
@@ -187,7 +188,7 @@ def shot_cleanup(
     def setup_physics_method_params(
         self,
         shot_id: int,
-        mds_conn: MDSConnection,
+        mds_conn: DataConnection,
         disruption_time: float,
         retrieval_settings: RetrievalSettings,
     ) -> PhysicsMethodParams:
@@ -198,8 +199,8 @@ def setup_physics_method_params(
         ----------
         shot_id : int
             The ID of the shot.
-        mds_conn : MDSConnection
-            The MDS connection for the shot.
+        mds_conn : DataConnection
+            The data connection for the shot.
         disruption_time : float
             The disruption time of the shot.
         retrieval_settings : RetrievalSettings
@@ -231,7 +232,7 @@ def setup_physics_method_params(
     def _init_times(
         self,
         shot_id: int,
-        mds_conn: MDSConnection,
+        mds_conn: DataConnection,
         disruption_time: float,
         retrieval_settings: RetrievalSettings,
     ) -> np.ndarray:
@@ -242,8 +243,8 @@ def _init_times(
         ----------
         shot_id : int
             The ID of the shot.
-        mds_conn : MDSConnection
-            The MDS connection for the shot.
+        mds_conn : DataConnection
+            The data connection for the shot.
         disruption_time : float
             The disruption time of the shot.
         retrieval_settings : RetrievalSettings

disruption_py/inout/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+"""Data connection abstractions and implementations."""
+
+from disruption_py.inout.base import DataConnection, ProcessConnection
+
+__all__ = ["DataConnection", "ProcessConnection"]

disruption_py/inout/base.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+
+"""
+Abstract base classes for data connections.
+
+DataConnection: per-shot data access (get_data, get_data_with_dims, get_dims).
+ProcessConnection: per-process factory that creates DataConnection instances.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Tuple
+
+import numpy as np
+
+from disruption_py.machine.tokamak import Tokamak
+
+
+class DataConnection(ABC):
+    """Per-shot data access interface.
+
+    Each instance is bound to a single shot. Implementations must provide
+    get_data, get_data_with_dims, get_dims, and cleanup. The reconnect
+    method is optional (default no-op).
+    """
+
+    @property
+    @abstractmethod
+    def shot_id(self) -> int:
+        """The shot ID this connection is bound to."""
+
+    @abstractmethod
+    def get_data(self, path: str, group: str = None, **kwargs) -> np.ndarray:
+        """Get data at path.
+
+        Parameters
+        ----------
+        path : str
+            Data path (node path for MDSplus, variable name for Xarray).
+        group : str, optional
+            Container name (tree for MDSplus, group for Xarray).
+        **kwargs
+            Backend-specific options.
+
+        Returns
+        -------
+        np.ndarray
+        """
+
+    @abstractmethod
+    def get_data_with_dims(
+        self,
+        path: str,
+        group: str = None,
+        dim_nums: List = None,
+        **kwargs,
+    ) -> Tuple:
+        """Get data and dimension arrays.
+
+        Parameters
+        ----------
+        path : str
+            Data path.
+        group : str, optional
+            Container name.
+        dim_nums : List, optional
+            Dimension indices to retrieve. Default [0].
+        **kwargs
+            Backend-specific options.
+
+        Returns
+        -------
+        Tuple
+            (data, dim0, dim1, ...) as numpy arrays.
+        """
+
+    @abstractmethod
+    def get_dims(
+        self,
+        path: str,
+        group: str = None,
+        dim_nums: List = None,
+        **kwargs,
+    ) -> Tuple:
+        """Get only dimension arrays.
+
+        Parameters
+        ----------
+        path : str
+            Data path.
+        group : str, optional
+            Container name.
+        dim_nums : List, optional
+            Dimension indices to retrieve. Default [0].
+        **kwargs
+            Backend-specific options.
+
+        Returns
+        -------
+        Tuple
+            Requested dimensions.
+        """
+
+    @abstractmethod
+    def cleanup(self) -> None:
+        """Release resources for this shot."""
+
+    def reconnect(self) -> None:
+        """Reconnect after error. Default no-op.
+
+        Xarray opens a new DataTree per shot so there is nothing to
+        reconnect. MDSplus overrides this to call conn.reconnect().
+        """
+
+
+class ProcessConnection(ABC):
+    """Per-process factory that creates DataConnection instances."""
+
+    @abstractmethod
+    def get_shot_connection(self, shot_id: int) -> DataConnection:
+        """Create a per-shot DataConnection for the given shot."""
+
+    @classmethod
+    @abstractmethod
+    def from_config(cls, tokamak: Tokamak) -> "ProcessConnection":
+        """Create a ProcessConnection from tokamak configuration."""