src/services/scpi_service.py - master - devsec/anothertap

src/services/scpi_service.py
333 строки13 KB
devsec
clean code gitignore
24 авг 2025, 21:27
24 авг 2025, 21:2775a9527
"""
SCPI Communication Service for controlling measurement devices
"""

import pyvisa
import time
import logging
from typing import Optional, Dict, Any, Tuple, TYPE_CHECKING
from threading import Lock
from contextlib import contextmanager

if TYPE_CHECKING:
    from pyvisa.resources import Resource

from models.device import Device, DeviceStatus, ConnectionType
from services.simulation_service import SimulationService


class SCPIService:
    """
    Service for managing SCPI communication with measurement devices
    """
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self._resource_manager: Optional[pyvisa.ResourceManager] = None
        self._connections: Dict[str, Any] = {}
        self._connection_lock = Lock()
        self.simulation_service = SimulationService()
        
    @property
    def resource_manager(self) -> pyvisa.ResourceManager:
        """Get or create VISA resource manager"""
        if self._resource_manager is None:
            try:
                self._resource_manager = pyvisa.ResourceManager()
            except Exception as e:
                self.logger.error(f"Failed to create VISA resource manager: {e}")
                # Fallback to PyVISA-py backend
                try:
                    self._resource_manager = pyvisa.ResourceManager('@py')
                except Exception as e2:
                    self.logger.error(f"Failed to create PyVISA-py resource manager: {e2}")
                    raise Exception("Could not initialize VISA resource manager")
        return self._resource_manager
    
    def list_available_resources(self) -> list[str]:
        """List all available VISA resources"""
        try:
            return list(self.resource_manager.list_resources())
        except Exception as e:
            self.logger.error(f"Failed to list resources: {e}")
            return []
    
    def connect_device(self, device: Device) -> bool:
        """
        Connect to a device and update its status
        
        Args:
            device: Device to connect to
            
        Returns:
            True if connection successful, False otherwise
        """
        with self._connection_lock:
            try:
                # Handle simulated devices
                if device.simulation_mode or device.connection_type == ConnectionType.SIMULATION:
                    return self._connect_simulated_device(device)
                
                if device.id in self._connections:
                    # Already connected
                    if self._test_connection(self._connections[device.id]):
                        device.status = DeviceStatus.CONNECTED
                        device.last_error = None
                        return True
                    else:
                        # Connection lost, remove it
                        self._disconnect_device_internal(device.id)
                
                # Create new connection
                resource = self.resource_manager.open_resource(device.connection_string)
                
                # Configure connection based on type
                self._configure_connection(resource, device)
                
                # Test connection with identification query
                response = resource.query("*IDN?", delay=0.1)
                self.logger.info(f"Connected to device {device.name}: {response.strip()}")
                
                # Store connection
                self._connections[device.id] = resource
                device.status = DeviceStatus.CONNECTED
                device.last_error = None
                
                # Update device info if possible
                self._update_device_info(device, response.strip())
                
                return True
                
            except Exception as e:
                error_msg = f"Failed to connect to device {device.name}: {str(e)}"
                self.logger.error(error_msg)
                device.status = DeviceStatus.ERROR
                device.last_error = error_msg
                return False
    
    def disconnect_device(self, device: Device) -> bool:
        """
        Disconnect from a device
        
        Args:
            device: Device to disconnect from
            
        Returns:
            True if disconnection successful, False otherwise
        """
        with self._connection_lock:
            # Handle simulated devices
            if device.simulation_mode or device.connection_type == ConnectionType.SIMULATION:
                self.simulation_service.remove_simulator(device.id)
                device.status = DeviceStatus.DISCONNECTED
                device.last_error = None
                return True
            
            success = self._disconnect_device_internal(device.id)
            if success:
                device.status = DeviceStatus.DISCONNECTED
                device.last_error = None
            return success
    
    def _disconnect_device_internal(self, device_id: str) -> bool:
        """Internal method to disconnect device without lock"""
        try:
            if device_id in self._connections:
                resource = self._connections[device_id]
                resource.close()
                del self._connections[device_id]
                self.logger.info(f"Disconnected from device {device_id}")
            return True
        except Exception as e:
            self.logger.error(f"Error disconnecting device {device_id}: {e}")
            return False
    
    def send_command(self, device: Device, command: str, timeout: float = 5.0) -> Tuple[bool, Optional[str]]:
        """
        Send SCPI command to device
        
        Args:
            device: Target device
            command: SCPI command to send
            timeout: Timeout in seconds
            
        Returns:
            Tuple of (success, response). Response is None for commands, string for queries
        """
        # Handle simulated devices
        if device.simulation_mode or device.connection_type == ConnectionType.SIMULATION:
            if device.status != DeviceStatus.CONNECTED:
                return False, None
            return self.simulation_service.send_command(device, command)
        
        # Handle real devices
        if device.id not in self._connections:
            if not self.connect_device(device):
                return False, None
        
        resource = self._connections[device.id]
        
        try:
            # Set timeout
            resource.timeout = int(timeout * 1000)  # Convert to milliseconds
            
            if command.strip().endswith('?'):
                # Query command - expects response
                response = resource.query(command)
                self.logger.debug(f"Query '{command}' -> '{response.strip()}'")
                return True, response.strip()
            else:
                # Write command - no response expected
                resource.write(command)
                self.logger.debug(f"Command sent: '{command}'")
                return True, None
                
        except pyvisa.VisaIOError as e:
            error_msg = f"VISA error sending command '{command}' to {device.name}: {e}"
            self.logger.error(error_msg)
            device.last_error = error_msg
            return False, None
        except Exception as e:
            error_msg = f"Error sending command '{command}' to {device.name}: {e}"
            self.logger.error(error_msg)
            device.last_error = error_msg
            return False, None
    
    def test_device_connection(self, device: Device) -> bool:
        """
        Test if device connection is working
        
        Args:
            device: Device to test
            
        Returns:
            True if connection is working, False otherwise
        """
        # Handle simulated devices
        if device.simulation_mode or device.connection_type == ConnectionType.SIMULATION:
            return device.status == DeviceStatus.CONNECTED
        
        # Handle real devices
        if device.id not in self._connections:
            return False
        
        return self._test_connection(self._connections[device.id])
    
    def _test_connection(self, resource: Any) -> bool:
        """Test a VISA resource connection"""
        try:
            resource.query("*IDN?", delay=0.1)
            return True
        except:
            return False
    
    def _configure_connection(self, resource: Any, device: Device):
        """Configure connection parameters based on device type"""
        try:
            if device.connection_type == ConnectionType.TCPIP:
                resource.timeout = 5000  # 5 seconds
                resource.read_termination = '\n'
                resource.write_termination = '\n'
            elif device.connection_type == ConnectionType.SERIAL:
                resource.timeout = 5000
                resource.baud_rate = 9600  # Default, should be configurable
                resource.data_bits = 8
                resource.parity = pyvisa.constants.Parity.none  # type: ignore
                resource.stop_bits = pyvisa.constants.StopBits.one  # type: ignore
                resource.read_termination = '\n'
                resource.write_termination = '\n'
            elif device.connection_type == ConnectionType.USB:
                resource.timeout = 5000
                resource.read_termination = '\n'
                resource.write_termination = '\n'
            elif device.connection_type == ConnectionType.GPIB:
                resource.timeout = 5000
                resource.read_termination = '\n'
                resource.write_termination = '\n'
                
        except Exception as e:
            self.logger.warning(f"Could not configure connection for {device.name}: {e}")
    
    def _update_device_info(self, device: Device, idn_response: str):
        """Update device information from *IDN? response"""
        try:
            # Parse IDN response (typically: manufacturer,model,serial,firmware)
            parts = [part.strip() for part in idn_response.split(',')]
            if len(parts) >= 2:
                if not device.manufacturer:
                    device.manufacturer = parts[0]
                if not device.model:
                    device.model = parts[1]
                if len(parts) >= 3 and not device.serial_number:
                    device.serial_number = parts[2]
                    
        except Exception as e:
            self.logger.debug(f"Could not parse IDN response: {e}")
    
    def _connect_simulated_device(self, device: Device) -> bool:
        """
        Connect to a simulated device
        
        Args:
            device: Device to simulate
            
        Returns:
            True if simulator created successfully
        """
        try:
            # Create simulator if it doesn't exist
            if self.simulation_service.create_simulator(device):
                device.status = DeviceStatus.CONNECTED
                device.last_error = None
                
                # Test simulator with *IDN? command
                success, response = self.simulation_service.send_command(device, "*IDN?")
                if success and response:
                    self.logger.info(f"Connected to simulated device {device.name}: {response}")
                    self._update_device_info(device, response)
                
                return True
            else:
                device.status = DeviceStatus.ERROR
                device.last_error = "Failed to create device simulator"
                return False
                
        except Exception as e:
            error_msg = f"Failed to connect to simulated device {device.name}: {str(e)}"
            self.logger.error(error_msg)
            device.status = DeviceStatus.ERROR
            device.last_error = error_msg
            return False
    
    @contextmanager
    def device_session(self, device: Device):
        """
        Context manager for device operations
        
        Usage:
            with scpi_service.device_session(device) as connected:
                if connected:
                    success, response = scpi_service.send_command(device, "*IDN?")
        """
        connected = self.connect_device(device)
        try:
            yield connected
        finally:
            # Optionally disconnect here, or keep connection alive
            pass
    
    def disconnect_all(self):
        """Disconnect from all devices"""
        with self._connection_lock:
            device_ids = list(self._connections.keys())
            for device_id in device_ids:
                self._disconnect_device_internal(device_id)
    
    def __del__(self):
        """Cleanup when service is destroyed"""
        try:
            self.disconnect_all()
            if self._resource_manager:
                self._resource_manager.close()
        except:
            pass