oscura 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

oscura/loaders/sigrok.py

@@ -60,90 +60,23 @@ def load_sigrok(
         sigrok session file format specification
     """
     path = Path(path)
-
-    if not path.exists():
-        raise LoaderError(
-            "File not found",
-            file_path=str(path),
-        )
-
-    if not zipfile.is_zipfile(path):
-        raise FormatError(
-            "File is not a valid sigrok session (not a ZIP archive)",
-            file_path=str(path),
-            expected="ZIP archive",
-        )
+    _validate_sigrok_file(path)
 
     try:
         with zipfile.ZipFile(path, "r") as zf:
-            # Parse metadata
+            # Parse metadata and extract channel info
             metadata_dict = _parse_metadata(zf, path)
+            sample_rate, channels, total_channels = _extract_channel_info(metadata_dict)
 
-            # Get sample rate from metadata
-            sample_rate = metadata_dict.get("samplerate", 1_000_000)
-
-            # Get channel information
-            channels = metadata_dict.get("channels", [])
-            total_channels = metadata_dict.get("total probes", len(channels))
-
-            # Find and read logic data files
-            logic_files = [name for name in zf.namelist() if name.startswith("logic-1")]
-
-            if not logic_files:
-                raise FormatError(
-                    "No logic data found in sigrok session",
-                    file_path=str(path),
-                    expected="logic-1-* data files",
-                )
-
-            # Read and combine logic data
-            data = _read_logic_data(zf, logic_files, total_channels)
-
-            # Select specific channel if requested
-            if channel is not None:
-                if isinstance(channel, int):
-                    if channel < 0 or channel >= data.shape[0]:
-                        raise LoaderError(
-                            f"Channel index {channel} out of range",
-                            file_path=str(path),
-                            details=f"Available channels: 0-{data.shape[0] - 1}",
-                        )
-                    channel_data = data[channel]
-                    channel_name = channels[channel] if channel < len(channels) else f"D{channel}"
-                elif isinstance(channel, str):
-                    if channel in channels:
-                        idx = channels.index(channel)
-                        channel_data = data[idx]
-                        channel_name = channel
-                    else:
-                        raise LoaderError(
-                            f"Channel '{channel}' not found",
-                            file_path=str(path),
-                            details=f"Available channels: {channels}",
-                        )
-                else:
-                    channel_data = data[0]  # type: ignore[unreachable]
-                    channel_name = channels[0] if channels else "D0"
-            else:
-                # Default to first channel
-                channel_data = data[0] if data.ndim > 1 else data
-                channel_name = channels[0] if channels else "D0"
-
-            # Compute edges
-            edges = _compute_edges(channel_data, sample_rate)
-
-            # Build metadata
-            trace_metadata = TraceMetadata(
-                sample_rate=float(sample_rate),
-                source_file=str(path),
-                channel_name=channel_name,
-                trigger_info=metadata_dict.get("trigger", None),
-            )
+            # Read logic data
+            data = _load_logic_data(zf, path, total_channels)
 
-            return DigitalTrace(
-                data=channel_data,
-                metadata=trace_metadata,
-                edges=edges,
+            # Select channel and build trace
+            channel_data, channel_name = _select_channel_data(data, channel, channels, path)
+
+            # Compute edges and build trace
+            return _build_digital_trace(
+                channel_data, channel_name, sample_rate, metadata_dict, path
             )
 
     except zipfile.BadZipFile as e:
@@ -163,6 +96,196 @@ def load_sigrok(
         ) from e
 
 
+def _validate_sigrok_file(path: Path) -> None:
+    """Validate that file exists and is a zip archive.
+
+    Args:
+        path: Path to sigrok file.
+
+    Raises:
+        LoaderError: If file not found.
+        FormatError: If not a valid ZIP file.
+    """
+    if not path.exists():
+        raise LoaderError("File not found", file_path=str(path))
+
+    if not zipfile.is_zipfile(path):
+        raise FormatError(
+            "File is not a valid sigrok session (not a ZIP archive)",
+            file_path=str(path),
+            expected="ZIP archive",
+        )
+
+
+def _extract_channel_info(metadata_dict: dict[str, Any]) -> tuple[float, list[str], int]:
+    """Extract channel information from metadata.
+
+    Args:
+        metadata_dict: Parsed metadata dictionary.
+
+    Returns:
+        Tuple of (sample_rate, channels, total_channels).
+    """
+    sample_rate = metadata_dict.get("samplerate", 1_000_000)
+    channels = metadata_dict.get("channels", [])
+    total_channels = metadata_dict.get("total probes", len(channels))
+    return float(sample_rate), channels, total_channels
+
+
+def _load_logic_data(zf: zipfile.ZipFile, path: Path, total_channels: int) -> NDArray[np.bool_]:
+    """Load logic data from sigrok session.
+
+    Args:
+        zf: Open ZipFile object.
+        path: Path to session file (for error messages).
+        total_channels: Total number of channels.
+
+    Returns:
+        Boolean array of shape (channels, samples).
+
+    Raises:
+        FormatError: If no logic data found.
+    """
+    logic_files = [name for name in zf.namelist() if name.startswith("logic-1")]
+
+    if not logic_files:
+        raise FormatError(
+            "No logic data found in sigrok session",
+            file_path=str(path),
+            expected="logic-1-* data files",
+        )
+
+    return _read_logic_data(zf, logic_files, total_channels)
+
+
+def _select_channel_data(
+    data: NDArray[np.bool_],
+    channel: str | int | None,
+    channels: list[str],
+    path: Path,
+) -> tuple[NDArray[np.bool_], str]:
+    """Select specific channel data or default to first channel.
+
+    Args:
+        data: Multi-channel data array.
+        channel: Channel selector (name, index, or None).
+        channels: List of channel names.
+        path: Path to file (for error messages).
+
+    Returns:
+        Tuple of (channel_data, channel_name).
+
+    Raises:
+        LoaderError: If channel not found or out of range.
+    """
+    if channel is None:
+        channel_data = data[0] if data.ndim > 1 else data
+        channel_name = channels[0] if channels else "D0"
+        return channel_data, channel_name
+
+    if isinstance(channel, int):
+        return _select_channel_by_index(data, channel, channels, path)
+    # isinstance(channel, str) must be true here
+    return _select_channel_by_name(data, channel, channels, path)
+
+
+def _select_channel_by_index(
+    data: NDArray[np.bool_],
+    channel: int,
+    channels: list[str],
+    path: Path,
+) -> tuple[NDArray[np.bool_], str]:
+    """Select channel by numeric index.
+
+    Args:
+        data: Multi-channel data array.
+        channel: Channel index.
+        channels: List of channel names.
+        path: Path to file (for error messages).
+
+    Returns:
+        Tuple of (channel_data, channel_name).
+
+    Raises:
+        LoaderError: If index out of range.
+    """
+    if channel < 0 or channel >= data.shape[0]:
+        raise LoaderError(
+            f"Channel index {channel} out of range",
+            file_path=str(path),
+            details=f"Available channels: 0-{data.shape[0] - 1}",
+        )
+    channel_data = data[channel]
+    channel_name = channels[channel] if channel < len(channels) else f"D{channel}"
+    return channel_data, channel_name
+
+
+def _select_channel_by_name(
+    data: NDArray[np.bool_],
+    channel: str,
+    channels: list[str],
+    path: Path,
+) -> tuple[NDArray[np.bool_], str]:
+    """Select channel by name.
+
+    Args:
+        data: Multi-channel data array.
+        channel: Channel name.
+        channels: List of channel names.
+        path: Path to file (for error messages).
+
+    Returns:
+        Tuple of (channel_data, channel_name).
+
+    Raises:
+        LoaderError: If channel name not found.
+    """
+    if channel in channels:
+        idx = channels.index(channel)
+        return data[idx], channel
+    else:
+        raise LoaderError(
+            f"Channel '{channel}' not found",
+            file_path=str(path),
+            details=f"Available channels: {channels}",
+        )
+
+
+def _build_digital_trace(
+    channel_data: NDArray[np.bool_],
+    channel_name: str,
+    sample_rate: float,
+    metadata_dict: dict[str, Any],
+    path: Path,
+) -> DigitalTrace:
+    """Build DigitalTrace object from channel data.
+
+    Args:
+        channel_data: Boolean array for selected channel.
+        channel_name: Name of the channel.
+        sample_rate: Sample rate in Hz.
+        metadata_dict: Metadata dictionary.
+        path: Path to source file.
+
+    Returns:
+        DigitalTrace object.
+    """
+    edges = _compute_edges(channel_data, sample_rate)
+
+    trace_metadata = TraceMetadata(
+        sample_rate=sample_rate,
+        source_file=str(path),
+        channel_name=channel_name,
+        trigger_info=metadata_dict.get("trigger"),
+    )
+
+    return DigitalTrace(
+        data=channel_data,
+        metadata=trace_metadata,
+        edges=edges,
+    )
+
+
 def _parse_metadata(zf: zipfile.ZipFile, path: Path) -> dict[str, Any]:
     """Parse sigrok session metadata.
 

oscura/loaders/tdms.py

@@ -16,6 +16,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
+from numpy.typing import NDArray
 
 from oscura.core.exceptions import FormatError, LoaderError
 from oscura.core.types import TraceMetadata, WaveformTrace
@@ -112,6 +113,17 @@ def _load_with_nptdms(
         FormatError: If file is not valid TDMS format or has no data.
         LoaderError: If channel or group not found.
     """
+    tdms_file = _parse_tdms_file(path)
+    target_group = _select_tdms_group(tdms_file, group, path)
+    target_channel = _select_tdms_channel(target_group, channel, path)
+    data = _extract_channel_data(target_channel, path)
+    metadata = _build_tdms_metadata(target_channel, target_group, tdms_file, path)
+
+    return WaveformTrace(data=data, metadata=metadata)
+
+
+def _parse_tdms_file(path: Path) -> Any:
+    """Parse TDMS file and validate structure."""
     try:
         tdms_file = TdmsFile.read(str(path))
     except Exception as e:
@@ -121,33 +133,37 @@ def _load_with_nptdms(
             expected="Valid NI TDMS format",
         ) from e
 
-    # Get available groups
     groups = list(tdms_file.groups())
-
     if not groups:
         raise FormatError(
             "No groups found in TDMS file",
             file_path=str(path),
         )
 
-    # Select group
-    if group is not None:
-        target_group = None
-        for g in groups:
-            if g.name == group:
-                target_group = g
-                break
-        if target_group is None:
-            available_groups = [g.name for g in groups]
-            raise LoaderError(
-                f"Group '{group}' not found",
-                file_path=str(path),
-                details=f"Available groups: {available_groups}",
-            )
-    else:
-        target_group = groups[0]
+    return tdms_file
+
+
+def _select_tdms_group(tdms_file: Any, group: str | None, path: Path) -> Any:
+    """Select target group from TDMS file."""
+    groups = list(tdms_file.groups())
+
+    if group is None:
+        return groups[0]
+
+    for g in groups:
+        if g.name == group:
+            return g
+
+    available_groups = [g.name for g in groups]
+    raise LoaderError(
+        f"Group '{group}' not found",
+        file_path=str(path),
+        details=f"Available groups: {available_groups}",
+    )
+
 
-    # Get channels in group
+def _select_tdms_channel(target_group: Any, channel: str | int | None, path: Path) -> Any:
+    """Select target channel from TDMS group."""
     channels = list(target_group.channels())
 
     if not channels:
@@ -156,35 +172,44 @@ def _load_with_nptdms(
             file_path=str(path),
         )
 
-    # Select channel
-    if channel is not None:
-        if isinstance(channel, int):
-            if channel < 0 or channel >= len(channels):
-                raise LoaderError(
-                    f"Channel index {channel} out of range",
-                    file_path=str(path),
-                    details=f"Available channels: 0-{len(channels) - 1}",
-                )
-            target_channel = channels[channel]
-        elif isinstance(channel, str):
-            target_channel = None
-            for ch in channels:
-                if ch.name == channel:
-                    target_channel = ch
-                    break
-            if target_channel is None:
-                available_channels = [ch.name for ch in channels]
-                raise LoaderError(
-                    f"Channel '{channel}' not found",
-                    file_path=str(path),
-                    details=f"Available channels: {available_channels}",
-                )
-        else:
-            target_channel = channels[0]  # type: ignore[unreachable]
+    if channel is None:
+        return channels[0]
+
+    if isinstance(channel, int):
+        return _select_channel_by_index(channels, channel, path)
+    elif isinstance(channel, str):
+        return _select_channel_by_name(channels, channel, path)
     else:
-        target_channel = channels[0]
+        return channels[0]  # type: ignore[unreachable]
+
 
-    # Get channel data
+def _select_channel_by_index(channels: list[Any], channel: int, path: Path) -> Any:
+    """Select channel by index."""
+    if channel < 0 or channel >= len(channels):
+        raise LoaderError(
+            f"Channel index {channel} out of range",
+            file_path=str(path),
+            details=f"Available channels: 0-{len(channels) - 1}",
+        )
+    return channels[channel]
+
+
+def _select_channel_by_name(channels: list[Any], channel: str, path: Path) -> Any:
+    """Select channel by name."""
+    for ch in channels:
+        if ch.name == channel:
+            return ch
+
+    available_channels = [ch.name for ch in channels]
+    raise LoaderError(
+        f"Channel '{channel}' not found",
+        file_path=str(path),
+        details=f"Available channels: {available_channels}",
+    )
+
+
+def _extract_channel_data(target_channel: Any, path: Path) -> NDArray[np.float64]:
+    """Extract and validate channel data."""
     data = target_channel.data
     if data is None or len(data) == 0:
         raise FormatError(
@@ -192,21 +217,21 @@ def _load_with_nptdms(
             file_path=str(path),
         )
 
-    # Convert to float64
-    data = np.asarray(data, dtype=np.float64)
+    return np.asarray(data, dtype=np.float64)
 
-    # Extract sample rate from properties
-    sample_rate = _get_sample_rate(target_channel, target_group, tdms_file)
 
-    # Extract other metadata
+def _build_tdms_metadata(
+    target_channel: Any,
+    target_group: Any,
+    tdms_file: Any,
+    path: Path,
+) -> TraceMetadata:
+    """Build metadata from TDMS channel properties."""
+    sample_rate = _get_sample_rate(target_channel, target_group, tdms_file)
     vertical_scale = target_channel.properties.get("NI_Scale[0]_Linear_Slope")
     vertical_offset = target_channel.properties.get("NI_Scale[0]_Linear_Y_Intercept")
 
-    # Get units if available
-    target_channel.properties.get("unit_string", None)
-
-    # Build metadata
-    metadata = TraceMetadata(
+    return TraceMetadata(
         sample_rate=sample_rate,
         vertical_scale=float(vertical_scale) if vertical_scale is not None else None,
         vertical_offset=float(vertical_offset) if vertical_offset is not None else None,
@@ -215,8 +240,6 @@ def _load_with_nptdms(
         trigger_info=_extract_tdms_properties(target_channel),
     )
 
-    return WaveformTrace(data=data, metadata=metadata)
-
 
 def _get_sample_rate(
     channel: Any,