<commit_before>from distributions import gaussian
from system import \
        Plant, RealChannel, IntegerChannel, BinarySymmetricChannel, LQGCost
import joint.control
import joint.coding
import separate.control.lloyd_max
import separate.control.noisy_lloyd_max
import separate.coding
import separate.coding.noisy_lloyd_max
import separate.coding.source
import separate.coding.convolutional
import trivial.coding
from utilities import memoized

from types import SimpleNamespace
import numpy as np
import inspect

class Simulation:
    def __init__(self, params):
        self.params = params

        self.plant = Plant(params.alpha, gaussian(params.W),
                gaussian(params.W), gaussian(params.V))
        self.LQG = LQGCost(self.plant, params.Q, params.R, params.F)

        # Globally known data
        self.globals = SimpleNamespace()
        # u[t]: control signal at time t
        self.globals.u = dict()
        if params.analog:
            # x_est_r[t1, t2]: receiver estimate of x(t1) at time t2
            self.globals.x_est_r = dict()

            self.channel = RealChannel(gaussian(1 / params.SNR))
        elif params.p == 0:
            self.channel = IntegerChannel(2**params.quantizer_bits)
        else:
            self.channel = BinarySymmetricChannel(params.p)

        if params.scheme == 'joint':
            self.observer = joint.control.Observer(self)
            self.controller = joint.control.Controller(self)
            if (params.KC, params.KS) == (1, 1):
                self.encoder = trivial.coding.Encoder(self)
                self.decoder = trivial.coding.Decoder(self)
            elif (params.KC, params.KS) == (2, 1):
                spiral_map = joint.coding.SpiralMap(params.omega, params.c_reg)
                self.encoder = joint.coding.Encoder(self, spiral_map)
                self.decoder = joint.coding.Decoder(self, spiral_map)

        elif params.scheme == 'lloyd-max':
            self.observer = separate.control.lloyd_max.Observer(self)
            self.controller = separate.control.lloyd_max.Controller(self)

            self.encoder = separate.coding.source.Encoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits))
            self.decoder = separate.coding.source.Decoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits))

        elif params.scheme == 'noisy-lloyd-max':
            self.observer = separate.control.noisy_lloyd_max.Observer(self)
            self.controller = separate.control.noisy_lloyd_max.Controller(self)

            self.convolutional_code = \
                    separate.coding.convolutional.ConvolutionalCode.random_code(
                    params.code_blocklength, params.quantizer_bits,
                    params.T - 1)
            self.encoder = separate.coding.noisy_lloyd_max.Encoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits), self.convolutional_code)
            self.decoder = separate.coding.noisy_lloyd_max.Decoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits), self.convolutional_code)

        elif params.scheme == 'separate':
            self.observer = separate.control.noisy_lloyd_max.Observer(self)
            self.controller = separate.control.noisy_lloyd_max.Controller(self)

            self.convolutional_code = \
                    separate.coding.convolutional.ConvolutionalCode.random_code(
                    params.code_blocklength, params.quantizer_bits,
                    params.T - 1)
            self.encoder = separate.coding.Encoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits), self.convolutional_code)
            self.decoder = separate.coding.Decoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits), self.convolutional_code)

    def simulate(self, T):
        self.t = 1
        while self.t <= T:
            # The observer observes the plant and generates a message
            msg = self.observer.observe(self.t, self.plant.y)
            # The encoder encodes the message
            code = self.encoder.encode(*msg)
            # The encoded message is sent over the channel
            code_recv = self.channel.transmit(code)
            # The decoder decodes the encoded message
            msg_recv = self.decoder.decode(*code_recv)
            # The controller receives the message and generates a control signal
            u = self.controller.control(self.t, *msg_recv)
            self.globals.u[self.t] = u

            yield self.t

            self.plant.step(u)
            self.LQG.step(u)

            self.t += 1
        yield self.t

    def measure_SDR(self):
        assert self.params.analog
        N_ITERATIONS = 10000
        signal_distr = gaussian(1)
        channel = self.channel.clone()
        distortions = []

        for signal in signal_distr.rvs(size=N_ITERATIONS):
            encoded = self.encoder.encode(signal)
            received = channel.transmit(encoded)
            (signal_estimate,) = self.decoder.decode(*received)
            distortions.append(signal - signal_estimate)

        return 1 / (np.array(distortions)**2).mean()

class Parameters:
    def __init__(self, T, alpha, W, V, Q, R, F):
        self.T = T # Time horizon
        self.alpha = alpha # System coefficient
        assert alpha > 1 # unstable

        # System noise power (variance) parameters
        self.W = W # V[w_t] = V[x_1]
        self.V = V # V[v_t]

        # LQG parameters
        self.Q = Q
        self.R = R
        self.F = F

        # Pre-evaluate L(t) in the right order to avoid blowing the stack
        for t in range(self.T, 0, -1):
            self.L(t)

    def setRates(self, KC, KS):
        if (KC, KS) not in [(1, 1), (2, 1)]:
            raise ValueError("{{KC = {}; KS = {}}} unsupported".format(KC, KS))
        else:
            self.KC = KC
            self.KS = KS

    def setBlocklength(self, code_blocklength):
        self.code_blocklength = code_blocklength

    def setAnalog(self, SNR):
        self.SNR = SNR # 1 / V[n_t]
        self.analog = True

    def setDigital(self, quantizer_bits, p=0):
        self.quantizer_bits = quantizer_bits
        self.analog = False
        self.p = p # Bit error probability

    def setScheme(self, scheme):
        self.scheme = scheme
        if scheme == 'joint':
            assert self.analog
            if (self.KC, self.KS) == (1, 1):
                # Channel code signal-distortion ratio, 1 / V[neff_t]
                self.SDR0 = self.SNR # Optimum is achieved
            elif (self.KC, self.KS) == (2, 1):
                # TODO Allow these parameters to be chosen
                self.c_reg = np.sqrt(2)
                self.omega = 1
                self.SDR0 = Simulation(self).measure_SDR()

        elif scheme == 'separate':
            assert self.analog
            pass # Nothing to do

        elif scheme == 'lloyd-max':
            assert not self.analog
            pass # Nothing to do

        elif scheme == 'noisy-lloyd-max':
            assert not self.analog
            pass # Nothing to do

        else:
            raise ValueError("Unrecognized scheme: {}".format(scheme))

    def all(self):
        return {k: v for k, v in inspect.getmembers(self)
                if not k.startswith('__')}

    # Statically known parameters computed recursively using memoization

    @memoized
    def Pt(self, t, t_obs):
        if (t, t_obs) == (1, 0):
            return self.W
        elif t_obs == t:
            return self.Kt(t) * self.V
        elif t_obs == t-1:
            return self.alpha**2 * self.Pt(t-1, t-1) + self.W
        else:
            raise ValueError("({}, {}) not on the form (t, t) or (t+1, t)"
                             .format(t, t_obs))

    @memoized
    def Kt(self, t):
        P = self.Pt(t, t-1)
        return P / (P + self.V)

    @memoized
    def Pr(self, t, t_obs):
        if (t, t_obs) == (1, 0):
            return self.W
        elif t_obs == t:
            return ((self.Pr(t, t-1) + self.SDR0 * self.Pt(t, t))
                    / (1 + self.SDR0))
        elif t_obs == t-1:
            return self.alpha**2 * self.Pr(t-1, t-1) + self.W
        else:
            raise ValueError("({}, {}) not on the form (t, t) or (t+1, t)"
                             .format(t, t_obs))

    @memoized
    def L(self, t):
        S = self.S(t+1)
        return self.alpha * S / (S + self.R)

    @memoized
    def S(self, t):
        if t == self.T + 1:
            return self.F
        elif 1 <= t <= self.T:
            S = self.S(t+1)
            return self.alpha**2 * self.R * S / (S + self.R) + self.Q
        else:
            raise ValueError(("S({}) undefined because {} is out of the range "
                    + "[1, T + 1 = {}]").format(t, t, self.T + 1))

    # Infinite horizon values
    @memoized
    def S_inf(self):
        return next(filter(lambda S: S > 0,
                np.polynomial.polynomial.Polynomial([
                    -self.Q * self.R,
                    -(self.Q + (self.alpha**2 - 1) * self.R),
                    1
                    ]).roots()))

    @memoized
    def Pt_inf(self):
        return next(filter(lambda S: S > 0,
                np.polynomial.polynomial.Polynomial([
                    -self.V * self.W,
                    -((self.alpha**2 - 1) * self.V + self.W),
                    1
                    ]).roots()))

    @memoized
    def Pt_inf_bar(self):
        Pt = self.Pt_inf()
        return Pt * self.V / (Pt + self.V)

    @memoized
    def LQR_inf_classical(self):
        Pt = self.Pt_inf()
        Pt_bar = self.Pt_inf_bar()
        return self.Q * Pt_bar + self.S_inf() * (Pt - Pt_bar)

    @memoized
    def LQR_inf_upper_bound(self):
        Pt = self.Pt_inf()
        Pt_bar = self.Pt_inf_bar()
        return (self.LQR_inf_classical()
                + (self.Q + (self.alpha**2 - 1) * self.S_inf())
                    / (1 + self.SDR0 - self.alpha**2)
                  * (Pt - Pt_bar))

    @memoized
    def LQR_inf_lower_bound(self):
        Pt = self.Pt_inf()
        Pt_bar = self.Pt_inf_bar()
        return (self.LQR_inf_classical()
                + (self.Q + (self.alpha**2 - 1) * self.S_inf())
                    / (1 + self.SDR_OPTA() - self.alpha**2)
                  * (Pt - Pt_bar))

    @memoized
    def SDR_OPTA(self):
        return (1 + self.SNR)**(self.KC / self.KS) - 1
<commit_msg>Added Params.text_description<commit_after>from distributions import gaussian
from system import \
        Plant, RealChannel, IntegerChannel, BinarySymmetricChannel, LQGCost
import joint.control
import joint.coding
import separate.control.lloyd_max
import separate.control.noisy_lloyd_max
import separate.coding
import separate.coding.noisy_lloyd_max
import separate.coding.source
import separate.coding.convolutional
import trivial.coding
from utilities import memoized

from types import SimpleNamespace
import numpy as np
import inspect

class Simulation:
    def __init__(self, params):
        self.params = params

        self.plant = Plant(params.alpha, gaussian(params.W),
                gaussian(params.W), gaussian(params.V))
        self.LQG = LQGCost(self.plant, params.Q, params.R, params.F)

        # Globally known data
        self.globals = SimpleNamespace()
        # u[t]: control signal at time t
        self.globals.u = dict()
        if params.analog:
            # x_est_r[t1, t2]: receiver estimate of x(t1) at time t2
            self.globals.x_est_r = dict()

            self.channel = RealChannel(gaussian(1 / params.SNR))
        elif params.p == 0:
            self.channel = IntegerChannel(2**params.quantizer_bits)
        else:
            self.channel = BinarySymmetricChannel(params.p)

        if params.scheme == 'joint':
            self.observer = joint.control.Observer(self)
            self.controller = joint.control.Controller(self)
            if (params.KC, params.KS) == (1, 1):
                self.encoder = trivial.coding.Encoder(self)
                self.decoder = trivial.coding.Decoder(self)
            elif (params.KC, params.KS) == (2, 1):
                spiral_map = joint.coding.SpiralMap(params.omega, params.c_reg)
                self.encoder = joint.coding.Encoder(self, spiral_map)
                self.decoder = joint.coding.Decoder(self, spiral_map)

        elif params.scheme == 'lloyd-max':
            self.observer = separate.control.lloyd_max.Observer(self)
            self.controller = separate.control.lloyd_max.Controller(self)

            self.encoder = separate.coding.source.Encoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits))
            self.decoder = separate.coding.source.Decoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits))

        elif params.scheme == 'noisy-lloyd-max':
            self.observer = separate.control.noisy_lloyd_max.Observer(self)
            self.controller = separate.control.noisy_lloyd_max.Controller(self)

            self.convolutional_code = \
                    separate.coding.convolutional.ConvolutionalCode.random_code(
                    params.code_blocklength, params.quantizer_bits,
                    params.T - 1)
            self.encoder = separate.coding.noisy_lloyd_max.Encoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits), self.convolutional_code)
            self.decoder = separate.coding.noisy_lloyd_max.Decoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits), self.convolutional_code)

        elif params.scheme == 'separate':
            self.observer = separate.control.noisy_lloyd_max.Observer(self)
            self.controller = separate.control.noisy_lloyd_max.Controller(self)

            self.convolutional_code = \
                    separate.coding.convolutional.ConvolutionalCode.random_code(
                    params.code_blocklength, params.quantizer_bits,
                    params.T - 1)
            self.encoder = separate.coding.Encoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits), self.convolutional_code)
            self.decoder = separate.coding.Decoder(
                    self, separate.coding.source.DistributionTracker(
                        self, 2**params.quantizer_bits), self.convolutional_code)

    def simulate(self, T):
        self.t = 1
        while self.t <= T:
            # The observer observes the plant and generates a message
            msg = self.observer.observe(self.t, self.plant.y)
            # The encoder encodes the message
            code = self.encoder.encode(*msg)
            # The encoded message is sent over the channel
            code_recv = self.channel.transmit(code)
            # The decoder decodes the encoded message
            msg_recv = self.decoder.decode(*code_recv)
            # The controller receives the message and generates a control signal
            u = self.controller.control(self.t, *msg_recv)
            self.globals.u[self.t] = u

            yield self.t

            self.plant.step(u)
            self.LQG.step(u)

            self.t += 1
        yield self.t

    def measure_SDR(self):
        assert self.params.analog
        N_ITERATIONS = 10000
        signal_distr = gaussian(1)
        channel = self.channel.clone()
        distortions = []

        for signal in signal_distr.rvs(size=N_ITERATIONS):
            encoded = self.encoder.encode(signal)
            received = channel.transmit(encoded)
            (signal_estimate,) = self.decoder.decode(*received)
            distortions.append(signal - signal_estimate)

        return 1 / (np.array(distortions)**2).mean()

class Parameters:
    def __init__(self, T, alpha, W, V, Q, R, F):
        self.T = T # Time horizon
        self.alpha = alpha # System coefficient
        assert alpha > 1 # unstable

        # System noise power (variance) parameters
        self.W = W # V[w_t] = V[x_1]
        self.V = V # V[v_t]

        # LQG parameters
        self.Q = Q
        self.R = R
        self.F = F

        # Pre-evaluate L(t) in the right order to avoid blowing the stack
        for t in range(self.T, 0, -1):
            self.L(t)

    def setRates(self, KC, KS):
        if (KC, KS) not in [(1, 1), (2, 1)]:
            raise ValueError("{{KC = {}; KS = {}}} unsupported".format(KC, KS))
        else:
            self.KC = KC
            self.KS = KS

    def setBlocklength(self, code_blocklength):
        self.code_blocklength = code_blocklength

    def setAnalog(self, SNR):
        self.SNR = SNR # 1 / V[n_t]
        self.analog = True

    def setDigital(self, quantizer_bits, p=0):
        self.quantizer_bits = quantizer_bits
        self.analog = False
        self.p = p # Bit error probability

    def setScheme(self, scheme):
        self.scheme = scheme
        if scheme == 'joint':
            assert self.analog
            if (self.KC, self.KS) == (1, 1):
                # Channel code signal-distortion ratio, 1 / V[neff_t]
                self.SDR0 = self.SNR # Optimum is achieved
            elif (self.KC, self.KS) == (2, 1):
                # TODO Allow these parameters to be chosen
                self.c_reg = np.sqrt(2)
                self.omega = 1
                self.SDR0 = Simulation(self).measure_SDR()

        elif scheme == 'separate':
            assert self.analog
            pass # Nothing to do

        elif scheme == 'lloyd-max':
            assert not self.analog
            pass # Nothing to do

        elif scheme == 'noisy-lloyd-max':
            assert not self.analog
            pass # Nothing to do

        else:
            raise ValueError("Unrecognized scheme: {}".format(scheme))

    def all(self):
        return {k: v for k, v in inspect.getmembers(self)
                if not k.startswith('__')}

    def text_description(self):
        return \
"""Plant: α = {}
Noise: W = {}, V = {}, SNR = {}
Cost: Q = {}, R = {}, F = {}
Channel uses per sample: {}""".format(self.alpha, self.W, self.V, self.SNR,
        self.Q, self.R, self.F, self.KC)

    # Statically known parameters computed recursively using memoization

    @memoized
    def Pt(self, t, t_obs):
        if (t, t_obs) == (1, 0):
            return self.W
        elif t_obs == t:
            return self.Kt(t) * self.V
        elif t_obs == t-1:
            return self.alpha**2 * self.Pt(t-1, t-1) + self.W
        else:
            raise ValueError("({}, {}) not on the form (t, t) or (t+1, t)"
                             .format(t, t_obs))

    @memoized
    def Kt(self, t):
        P = self.Pt(t, t-1)
        return P / (P + self.V)

    @memoized
    def Pr(self, t, t_obs):
        if (t, t_obs) == (1, 0):
            return self.W
        elif t_obs == t:
            return ((self.Pr(t, t-1) + self.SDR0 * self.Pt(t, t))
                    / (1 + self.SDR0))
        elif t_obs == t-1:
            return self.alpha**2 * self.Pr(t-1, t-1) + self.W
        else:
            raise ValueError("({}, {}) not on the form (t, t) or (t+1, t)"
                             .format(t, t_obs))

    @memoized
    def L(self, t):
        S = self.S(t+1)
        return self.alpha * S / (S + self.R)

    @memoized
    def S(self, t):
        if t == self.T + 1:
            return self.F
        elif 1 <= t <= self.T:
            S = self.S(t+1)
            return self.alpha**2 * self.R * S / (S + self.R) + self.Q
        else:
            raise ValueError(("S({}) undefined because {} is out of the range "
                    + "[1, T + 1 = {}]").format(t, t, self.T + 1))

    # Infinite horizon values
    @memoized
    def S_inf(self):
        return next(filter(lambda S: S > 0,
                np.polynomial.polynomial.Polynomial([
                    -self.Q * self.R,
                    -(self.Q + (self.alpha**2 - 1) * self.R),
                    1
                    ]).roots()))

    @memoized
    def Pt_inf(self):
        return next(filter(lambda S: S > 0,
                np.polynomial.polynomial.Polynomial([
                    -self.V * self.W,
                    -((self.alpha**2 - 1) * self.V + self.W),
                    1
                    ]).roots()))

    @memoized
    def Pt_inf_bar(self):
        Pt = self.Pt_inf()
        return Pt * self.V / (Pt + self.V)

    @memoized
    def LQR_inf_classical(self):
        Pt = self.Pt_inf()
        Pt_bar = self.Pt_inf_bar()
        return self.Q * Pt_bar + self.S_inf() * (Pt - Pt_bar)

    @memoized
    def LQR_inf_upper_bound(self):
        Pt = self.Pt_inf()
        Pt_bar = self.Pt_inf_bar()
        return (self.LQR_inf_classical()
                + (self.Q + (self.alpha**2 - 1) * self.S_inf())
                    / (1 + self.SDR0 - self.alpha**2)
                  * (Pt - Pt_bar))

    @memoized
    def LQR_inf_lower_bound(self):
        Pt = self.Pt_inf()
        Pt_bar = self.Pt_inf_bar()
        return (self.LQR_inf_classical()
                + (self.Q + (self.alpha**2 - 1) * self.S_inf())
                    / (1 + self.SDR_OPTA() - self.alpha**2)
                  * (Pt - Pt_bar))

    @memoized
    def SDR_OPTA(self):
        return (1 + self.SNR)**(self.KC / self.KS) - 1
<|endoftext|>


<commit_before>import logging
log = logging.getLogger(__name__)

from threading import RLock

from spacq.tool.box import Enum, Synchronized

"""
Hardware device abstraction interface.
"""

# PyVISA, Linux GPIB, PyVISA USB.
drivers = Enum(['pyvisa', 'lgpib', 'pyvisa_usb'])

# Try to import all available drivers.
available_drivers = []

try:
	import Gpib
	import gpib
except ImportError:
	pass
else:
	available_drivers.append(drivers.lgpib)

try:
	import visa
except ImportError:
	pass
else:
	available_drivers.append(drivers.pyvisa)
	available_drivers.append(drivers.pyvisa_usb)

class DeviceNotFoundError(Exception):
	"""
	Failure to connect to a device.
	"""

	pass

class IbstaBits(object):
	"""
	Status bits (in ibsta) as reported by Linux GPIB.
	"""

	DCAS = 0x1
	DTAS = 0x2
	LACS = 0x4
	TACS = 0x8
	ATN = 0x10
	CIC = 0x20
	REM = 0x40
	LOK = 0x80
	CMPL = 0x100
	EVENT = 0x200
	SPOLL = 0x400
	RQS = 0x800,
	SRQI = 0x1000
	END = 0x2000
	TIMO = 0x4000
	ERR = 0x8000

class USBDevice(visa.Instrument):
	"""
	Using USB devices with PyVISA requires a small hack: the object must be an Instrument, but we can't call Instrument.__init__.
	"""

	def __init__(self, *args, **kwargs):
		# Bypass the initialization in visa.Instrument, due to "send_end" not being valid for USB.
		visa.ResourceTemplate.__init__(self, *args, **kwargs)

class SuperDevice(object):
	def _setup(self):
		"""
		Pre-connection setup.
		"""

		self.name = self.__class__.__name__

		self.resources = {}
		self.subdevices = {}

	def _connected(self):
		"""
		Post-connection setup.
		"""

		if hasattr(self, 'driver') and self.driver == drivers.lgpib:
			# Some devices don't assert the EOI line, so look for their EOS character instead.
			if hasattr(self, 'eos_char'):
				self.device.config(gpib.IbcEOSchar, ord(self.eos_char))
				self.device.config(gpib.IbcEOSrd, 1)

			# Gpib.Gpib doesn't complain if the device at the PAD doesn't actually exist.
			try:
				log.debug('GPIB device IDN: {0!r}'.format(self.idn))
			except gpib.GpibError as e:
				raise DeviceNotFoundError('Could not open device at "{0}".'.format(self.connection_resource), e)

		# Recursively.
		for name, subdev in self.subdevices.items():
			log.debug('Post-connection for subdevice "{0}".'.format(name))

			subdev._connected()

class AbstractDevice(SuperDevice):
	"""
	A class for controlling devices which can be connected to either via Ethernet and PyVISA or GPIB and Linux GPIB.
	"""

	def _setup(self):
		self.multi_command = None
		self.responses_expected = 0

		SuperDevice._setup(self)

		self.lock = RLock()

		self.status = []

	def __init__(self, ip_address=None, gpib_board=0, gpib_pad=None, gpib_sad=0,
			usb_resource=None, autoconnect=True):
		"""
		Ethernet (tcpip::<ip_address>::instr):
			ip_address: Address on which the device is listening on port 111.

		GPIB (gpib[gpib_board]::<gpib_pad>[::<gpib_sad>]::instr):
			gpib_board: GPIB board index. Defaults to 0.
			gpib_pad: Primary address of the device.
			gpib_sad: Secondary address of the device. Defaults to 0.

		USB (usb_resource):
			usb_resource: VISA resource of the form: USB[board]::<vendor>::<product>::<serial>[::<interface>]::RAW

		autoconnect: Connect to the device upon instantiation.
		"""

		self._setup()

		log.info('Creating device "{0}".'.format(self.name))

		if ip_address is not None:
			if drivers.pyvisa in available_drivers:
				log.debug('Using PyVISA with ip_address="{0}".'.format(ip_address))
				self.driver = drivers.pyvisa
				self.connection_resource = {
					'resource_name': 'tcpip::{0}::instr'.format(ip_address),
				}
			else:
				raise NotImplementedError('PyVISA required, but not available.')
		elif gpib_pad is not None:
			if drivers.lgpib in available_drivers:
				log.debug('Using Linux GPIB with gpib_board="{0}", gpib_pad="{1}", '
						'gpib_sad="{2}".'.format(gpib_board, gpib_pad, gpib_sad))
				self.driver = drivers.lgpib
				self.connection_resource = {
					'name': gpib_board,
					'pad': gpib_pad,
					'sad': gpib_sad,
				}
			elif drivers.pyvisa in available_drivers:
				log.debug('Using PyVISA with gpib_board="{0}", gpib_pad="{1}", '
						'gpib_sad="{2}".'.format(gpib_board, gpib_pad, gpib_sad))
				self.driver = drivers.pyvisa
				self.connection_resource = {
					'resource_name': 'gpib{0}::{1}::{2}::instr'.format(gpib_board, gpib_pad, gpib_sad),
				}
			else:
				raise NotImplementedError('Linux GPIB or PyVISA required, but not available.')
		elif usb_resource is not None:
			if drivers.pyvisa_usb in available_drivers:
				log.debug('Using PyVISA with usb_resource="{0}"'.format(usb_resource))
				self.driver = drivers.pyvisa_usb
				self.connection_resource = {
					'resource_name': usb_resource,
				}
			else:
				raise NotImplementedError('PyVISA required, but not available.')
		else:
			raise ValueError('Either an IP, GPIB, or USB address must be specified.')

		if autoconnect:
			self.connect()

	def __repr__(self):
		return '<{0}>'.format(self.__class__.__name__)

	def connect(self):
		"""
		Make a connection to the device.
		"""

		log.info('Connecting to device "{0}" using {1} at "{2}".'.format(self.name, self.driver, self.connection_resource))

		if self.driver == drivers.pyvisa:
			try:
				self.device = visa.Instrument(**self.connection_resource)
			except visa.VisaIOError as e:
				raise DeviceNotFoundError('Could not open device at "{0}".'.format(self.connection_resource), e)
		elif self.driver == drivers.lgpib:
			try:
				self.device = Gpib.Gpib(**self.connection_resource)
			except gpib.GpibError as e:
				raise DeviceNotFoundError('Could not open device at "{0}".'.format(self.connection_resource), e)
		elif self.driver == drivers.pyvisa_usb:
			try:
				self.device = USBDevice(**self.connection_resource)
			except visa.VisaIOError as e:
				raise DeviceNotFoundError('Could not open device at "{0}".'.format(self.connection_resource), e)

		try:
			self._connected()
		except Exception as e:
			raise DeviceNotFoundError('Could finish connection to device at "{0}".'.format(self.connection_resource), e)

	def multi_command_start(self):
		"""
		Redirect further commands to a buffer.
		"""

		log.debug('Starting multi-command message for device "{0}"'.format(self.name))

		if self.driver not in [drivers.pyvisa, drivers.lgpib]:
			raise NotImplementedError('Unsupported driver: "{0}".'.format(self.driver))

		self.multi_command = []
		self.responses_expected = 0

	@Synchronized()
	def multi_command_stop(self):
		"""
		Stop redirecting to a buffer, and send the buffered commands.

		Returns the results of queries if any were expected.
		"""

		log.debug('Stopping multi-command message for device "{0}"'.format(self.name))

		if self.multi_command is None:
			raise ValueError('Multi-command message not started.')
		elif not self.multi_command:
			# No commands.
			return []

		commands = self.multi_command
		# This ensures that write and ask will not buffer the real message.
		self.multi_command = None

		# Only commands not starting with "*" get a ":" prefix.
		commands = [cmd if cmd[0] == '*' else ':' + cmd for cmd in commands]
		message = ';'.join(commands)

		if self.responses_expected:
			result = self.ask(message)

			# FIXME: What if the response contains a meaningful ";" somewhere?
			return result.split(';', self.responses_expected - 1)
		else:
			self.write(message)

			return []

	@Synchronized()
	def write(self, message):
		"""
		Write to the device.

		Supports multi-command.
		"""

		if self.multi_command is not None:
			log.debug('Writing to multi-command buffer for device "{0}": {1!r}'.format(self.name, message))

			self.multi_command.append(message)
			return

		log.debug('Writing to device "{0}": {1!r}'.format(self.name, message))

		if self.driver in [drivers.pyvisa, drivers.lgpib]:
			self.device.write(message)
		elif self.driver == drivers.pyvisa_usb:
			# Send the message raw.
			visa.vpp43.write(self.device.vi, message)

	@Synchronized()
	def read_raw(self, chunk_size=512):
		"""
		Read everything the device has to say and return it exactly.
		"""

		log.debug('Reading from device "{0}".'.format(self.name))

		buf = ''

		if self.driver in [drivers.pyvisa, drivers.pyvisa_usb]:
			buf = self.device.read_raw()
		elif self.driver == drivers.lgpib:
			status = 0
			while not status:
				buf += self.device.read(len=chunk_size)
				status = self.device.ibsta() & IbstaBits.END

		log.debug('Read from device "{0}": {1!r}'.format(self.name, buf))

		return buf

	def read(self):
		"""
		Read from the device, but strip terminating whitespace.
		"""

		return self.read_raw().rstrip()

	@Synchronized()
	def ask_raw(self, message):
		"""
		Write, then read_raw.
		"""

		self.write(message)
		return self.read_raw()

	@Synchronized()
	def ask(self, message):
		"""
		Write, then read.

		Supports multi-command.
		"""

		self.write(message)

		if self.multi_command is None:
			return self.read()
		else:
			self.responses_expected += 1

	def close(self):
		"""
		Close the connection, if possible.
		"""

		log.debug('Closing device: {0}'.format(self.name))

		if self.driver in [drivers.pyvisa, drivers.pyvisa_usb]:
			self.device.close()

	def find_resource(self, path):
		"""
		Return a Resource given a resource path spec.

		eg. ('subdevice A', 'subdevice B', 'resource C') -> Resource
		"""

		log.debug('Looking for resource {0}.'.format(path))

		if len(path) < 1:
			raise ValueError('No path provided.')

		# Keep track of the last device in the tree so far.
		dev = self
		# Keep track of the path so far.
		traversed = ()

		while len(path) > 1:
			try:
				dev = dev.subdevices[path[0]]
			except KeyError:
				raise ValueError('No subdevice "{0}" in {1}.'.format(path[0], traversed))

			traversed += path[:1]
			path = path[1:]

		try:
			return dev.resources[path[0]]
		except KeyError:
			raise ValueError('No resource "{0}" in {1}.'.format(path[0], traversed))

	@property
	def idn(self):
		"""
		Ask the device for identification.
		"""

		if self.driver in [drivers.pyvisa, drivers.lgpib]:
			return self.ask('*idn?')

	@property
	def opc(self):
		"""
		Wait until the device is done.
		"""

		while True:
			if self.driver == drivers.pyvisa:
				try:
					self.ask('*opc?')
				except visa.VisaIOError as e:
					if e.error_code != visa.VI_ERROR_TMO:
						raise
				else:
					break
			elif self.driver == drivers.lgpib:
				try:
					self.ask('*opc?')
				except gpib.GpibError as e:
					# TODO: Keep going on timeout.
					raise
				else:
					break

class AbstractSubdevice(SuperDevice):
	"""
	A subdevice (eg. channel) of a hardware device.
	"""

	def _setup(self):
		SuperDevice._setup(self)

		# Synchronized methods should use the device lock.
		self.lock = self.device.lock if self.device else None

	def __init__(self, device):
		self.device = device

		self._setup()
<commit_msg>Added driver-agnostic DeviceTimeout Exception class<commit_after>import logging
log = logging.getLogger(__name__)

from threading import RLock
from time import time

from spacq.tool.box import Enum, Synchronized

"""
Hardware device abstraction interface.
"""

# PyVISA, Linux GPIB, PyVISA USB.
drivers = Enum(['pyvisa', 'lgpib', 'pyvisa_usb'])

# Try to import all available drivers.
available_drivers = []

try:
	import Gpib
	import gpib
except ImportError:
	pass
else:
	available_drivers.append(drivers.lgpib)

try:
	import visa
except ImportError:
	pass
else:
	available_drivers.append(drivers.pyvisa)
	available_drivers.append(drivers.pyvisa_usb)

class DeviceNotFoundError(Exception):
	"""
	Failure to connect to a device.
	"""

	pass

class DeviceTimeout(Exception):
	"""
	Timeout on action.
	"""

	pass

class IbstaBits(object):
	"""
	Status bits (in ibsta) as reported by Linux GPIB.
	"""

	DCAS = 0x1
	DTAS = 0x2
	LACS = 0x4
	TACS = 0x8
	ATN = 0x10
	CIC = 0x20
	REM = 0x40
	LOK = 0x80
	CMPL = 0x100
	EVENT = 0x200
	SPOLL = 0x400
	RQS = 0x800,
	SRQI = 0x1000
	END = 0x2000
	TIMO = 0x4000
	ERR = 0x8000

class USBDevice(visa.Instrument):
	"""
	Using USB devices with PyVISA requires a small hack: the object must be an Instrument, but we can't call Instrument.__init__.
	"""

	def __init__(self, *args, **kwargs):
		# Bypass the initialization in visa.Instrument, due to "send_end" not being valid for USB.
		visa.ResourceTemplate.__init__(self, *args, **kwargs)

class SuperDevice(object):
	def _setup(self):
		"""
		Pre-connection setup.
		"""

		self.name = self.__class__.__name__

		self.resources = {}
		self.subdevices = {}

	def _connected(self):
		"""
		Post-connection setup.
		"""

		if hasattr(self, 'driver') and self.driver == drivers.lgpib:
			# Some devices don't assert the EOI line, so look for their EOS character instead.
			if hasattr(self, 'eos_char'):
				self.device.config(gpib.IbcEOSchar, ord(self.eos_char))
				self.device.config(gpib.IbcEOSrd, 1)

			# Gpib.Gpib doesn't complain if the device at the PAD doesn't actually exist.
			try:
				log.debug('GPIB device IDN: {0!r}'.format(self.idn))
			except gpib.GpibError as e:
				raise DeviceNotFoundError('Could not open device at "{0}".'.format(self.connection_resource), e)

		# Recursively.
		for name, subdev in self.subdevices.items():
			log.debug('Post-connection for subdevice "{0}".'.format(name))

			subdev._connected()

class AbstractDevice(SuperDevice):
	"""
	A class for controlling devices which can be connected to either via Ethernet and PyVISA or GPIB and Linux GPIB.
	"""

	max_timeout = 15 # s

	def _setup(self):
		self.multi_command = None
		self.responses_expected = 0

		SuperDevice._setup(self)

		self.lock = RLock()

		self.status = []

	def __init__(self, ip_address=None, gpib_board=0, gpib_pad=None, gpib_sad=0,
			usb_resource=None, autoconnect=True):
		"""
		Ethernet (tcpip::<ip_address>::instr):
			ip_address: Address on which the device is listening on port 111.

		GPIB (gpib[gpib_board]::<gpib_pad>[::<gpib_sad>]::instr):
			gpib_board: GPIB board index. Defaults to 0.
			gpib_pad: Primary address of the device.
			gpib_sad: Secondary address of the device. Defaults to 0.

		USB (usb_resource):
			usb_resource: VISA resource of the form: USB[board]::<vendor>::<product>::<serial>[::<interface>]::RAW

		autoconnect: Connect to the device upon instantiation.
		"""

		self._setup()

		log.info('Creating device "{0}".'.format(self.name))

		if ip_address is not None:
			if drivers.pyvisa in available_drivers:
				log.debug('Using PyVISA with ip_address="{0}".'.format(ip_address))
				self.driver = drivers.pyvisa
				self.connection_resource = {
					'resource_name': 'tcpip::{0}::instr'.format(ip_address),
				}
			else:
				raise NotImplementedError('PyVISA required, but not available.')
		elif gpib_pad is not None:
			if drivers.lgpib in available_drivers:
				log.debug('Using Linux GPIB with gpib_board="{0}", gpib_pad="{1}", '
						'gpib_sad="{2}".'.format(gpib_board, gpib_pad, gpib_sad))
				self.driver = drivers.lgpib
				self.connection_resource = {
					'name': gpib_board,
					'pad': gpib_pad,
					'sad': gpib_sad,
				}
			elif drivers.pyvisa in available_drivers:
				log.debug('Using PyVISA with gpib_board="{0}", gpib_pad="{1}", '
						'gpib_sad="{2}".'.format(gpib_board, gpib_pad, gpib_sad))
				self.driver = drivers.pyvisa
				self.connection_resource = {
					'resource_name': 'gpib{0}::{1}::{2}::instr'.format(gpib_board, gpib_pad, gpib_sad),
				}
			else:
				raise NotImplementedError('Linux GPIB or PyVISA required, but not available.')
		elif usb_resource is not None:
			if drivers.pyvisa_usb in available_drivers:
				log.debug('Using PyVISA with usb_resource="{0}"'.format(usb_resource))
				self.driver = drivers.pyvisa_usb
				self.connection_resource = {
					'resource_name': usb_resource,
				}
			else:
				raise NotImplementedError('PyVISA required, but not available.')
		else:
			raise ValueError('Either an IP, GPIB, or USB address must be specified.')

		if autoconnect:
			self.connect()

	def __repr__(self):
		return '<{0}>'.format(self.__class__.__name__)

	def connect(self):
		"""
		Make a connection to the device.
		"""

		log.info('Connecting to device "{0}" using {1} at "{2}".'.format(self.name, self.driver, self.connection_resource))

		if self.driver == drivers.pyvisa:
			try:
				self.device = visa.Instrument(**self.connection_resource)
			except visa.VisaIOError as e:
				raise DeviceNotFoundError('Could not open device at "{0}".'.format(self.connection_resource), e)
		elif self.driver == drivers.lgpib:
			try:
				self.device = Gpib.Gpib(**self.connection_resource)
			except gpib.GpibError as e:
				raise DeviceNotFoundError('Could not open device at "{0}".'.format(self.connection_resource), e)
		elif self.driver == drivers.pyvisa_usb:
			try:
				self.device = USBDevice(**self.connection_resource)
			except visa.VisaIOError as e:
				raise DeviceNotFoundError('Could not open device at "{0}".'.format(self.connection_resource), e)

		try:
			self._connected()
		except Exception as e:
			raise DeviceNotFoundError('Could finish connection to device at "{0}".'.format(self.connection_resource), e)

	def multi_command_start(self):
		"""
		Redirect further commands to a buffer.
		"""

		log.debug('Starting multi-command message for device "{0}"'.format(self.name))

		if self.driver not in [drivers.pyvisa, drivers.lgpib]:
			raise NotImplementedError('Unsupported driver: "{0}".'.format(self.driver))

		self.multi_command = []
		self.responses_expected = 0

	@Synchronized()
	def multi_command_stop(self):
		"""
		Stop redirecting to a buffer, and send the buffered commands.

		Returns the results of queries if any were expected.
		"""

		log.debug('Stopping multi-command message for device "{0}"'.format(self.name))

		if self.multi_command is None:
			raise ValueError('Multi-command message not started.')
		elif not self.multi_command:
			# No commands.
			return []

		commands = self.multi_command
		# This ensures that write and ask will not buffer the real message.
		self.multi_command = None

		# Only commands not starting with "*" get a ":" prefix.
		commands = [cmd if cmd[0] == '*' else ':' + cmd for cmd in commands]
		message = ';'.join(commands)

		if self.responses_expected:
			result = self.ask(message)

			# FIXME: What if the response contains a meaningful ";" somewhere?
			return result.split(';', self.responses_expected - 1)
		else:
			self.write(message)

			return []

	@Synchronized()
	def write(self, message):
		"""
		Write to the device.

		Supports multi-command.
		"""

		if self.multi_command is not None:
			log.debug('Writing to multi-command buffer for device "{0}": {1!r}'.format(self.name, message))

			self.multi_command.append(message)
			return

		log.debug('Writing to device "{0}": {1!r}'.format(self.name, message))

		if self.driver == drivers.pyvisa:
			try:
				self.device.write(message)
			except visa.VisaIOError as e:
				if e.error_code == visa.VI_ERROR_TMO:
					raise DeviceTimeout(e)
				else:
					raise
		elif self.driver == drivers.lgpib:
			try:
				self.device.write(message)
			except gpib.GpibError as e:
				if 'timeout' in e.msg:
					raise DeviceTimeout(e)
				else:
					raise
		elif self.driver == drivers.pyvisa_usb:
			# Send the message raw.
			visa.vpp43.write(self.device.vi, message)

	@Synchronized()
	def read_raw(self, chunk_size=512):
		"""
		Read everything the device has to say and return it exactly.
		"""

		log.debug('Reading from device "{0}".'.format(self.name))

		buf = ''

		if self.driver in [drivers.pyvisa, drivers.pyvisa_usb]:
			try:
				buf = self.device.read_raw()
			except visa.VisaIOError as e:
				if e.error_code == visa.VI_ERROR_TMO:
					raise DeviceTimeout(e)
				else:
					raise
		elif self.driver == drivers.lgpib:
			status = 0
			while status == 0:
				try:
					buf += self.device.read(len=chunk_size)
				except gpib.GpibError as e:
					if 'timeout' in e.msg:
						raise DeviceTimeout(e)
					else:
						raise

				status = self.device.ibsta() & IbstaBits.END

		log.debug('Read from device "{0}": {1!r}'.format(self.name, buf))

		return buf

	def read(self):
		"""
		Read from the device, but strip terminating whitespace.
		"""

		return self.read_raw().rstrip()

	@Synchronized()
	def ask_raw(self, message):
		"""
		Write, then read_raw.
		"""

		self.write(message)
		return self.read_raw()

	@Synchronized()
	def ask(self, message):
		"""
		Write, then read.

		Supports multi-command.
		"""

		self.write(message)

		if self.multi_command is None:
			return self.read()
		else:
			self.responses_expected += 1

	def close(self):
		"""
		Close the connection, if possible.
		"""

		log.debug('Closing device: {0}'.format(self.name))

		if self.driver in [drivers.pyvisa, drivers.pyvisa_usb]:
			self.device.close()

	def find_resource(self, path):
		"""
		Return a Resource given a resource path spec.

		eg. ('subdevice A', 'subdevice B', 'resource C') -> Resource
		"""

		log.debug('Looking for resource {0}.'.format(path))

		if len(path) < 1:
			raise ValueError('No path provided.')

		# Keep track of the last device in the tree so far.
		dev = self
		# Keep track of the path so far.
		traversed = ()

		while len(path) > 1:
			try:
				dev = dev.subdevices[path[0]]
			except KeyError:
				raise ValueError('No subdevice "{0}" in {1}.'.format(path[0], traversed))

			traversed += path[:1]
			path = path[1:]

		try:
			return dev.resources[path[0]]
		except KeyError:
			raise ValueError('No resource "{0}" in {1}.'.format(path[0], traversed))

	@property
	def idn(self):
		"""
		Ask the device for identification.
		"""

		if self.driver in [drivers.pyvisa, drivers.lgpib]:
			return self.ask('*idn?')

	@property
	def opc(self):
		"""
		Wait until the device is done.
		"""

		end_time = time() + self.max_timeout

		while True:
			if self.driver in [drivers.pyvisa, drivers.lgpib]:
				try:
					self.ask('*opc?')
				except DeviceTimeout:
					if time() > end_time:
						raise
				else:
					break

class AbstractSubdevice(SuperDevice):
	"""
	A subdevice (eg. channel) of a hardware device.
	"""

	def _setup(self):
		SuperDevice._setup(self)

		# Synchronized methods should use the device lock.
		self.lock = self.device.lock if self.device else None

	def __init__(self, device):
		self.device = device

		self._setup()
<|endoftext|>


<commit_before>__version__ = '1.9.3.3'
<commit_msg>1.9.4<commit_after>__version__ = '1.9.4'
<|endoftext|>


<commit_before>"""Copy number detection using read counts, with cn.mops.

http://www.bioconductor.org/packages/release/bioc/html/cn.mops.html
"""
from contextlib import closing
import os
import shutil
import subprocess

import pysam

from bcbio import bam, install, utils
from bcbio.distributed.transaction import file_transaction
from bcbio.log import logger
from bcbio.pipeline import config_utils
from bcbio.provenance import do
from bcbio.variation import vcfutils

def run(items):
    """Detect copy number variations from batched set of samples using cn.mops.
    """
    names = [x["name"][-1] for x in items]
    work_bams = [x["align_bam"] for x in items]
    if len(items) == 1:
        raise ValueError("cn.mops only works on batches with multiple samples")
    data = items[0]
    # XXX Can parallelize cn.mops with snow but currently causing
    # errors. Could also parallelize by chromosome.
    #num_cores = data["config"]["algorithm"].get("num_cores", 1)
    work_dir = utils.safe_makedir(os.path.join(data["dirs"]["work"], "structural", names[0],
                                               "cn_mops"))
    with closing(pysam.Samfile(work_bams[0], "rb")) as pysam_work_bam:
        out_files = [_run_on_chrom(chrom, work_bams, names, work_dir, items)
                     for chrom in pysam_work_bam.references]
    out_file = _combine_out_files(out_files, work_bams[0], work_dir)
    out = []
    for data in items:
        if "sv" not in data:
            data["sv"] = {}
        data["sv"]["cnv"] = _prep_sample_cnvs(out_file, data)
        out.append(data)
    return out

def _combine_out_files(chr_files, base_bam, work_dir):
    """Concatenate all CNV calls into a single file.
    """
    out_file = os.path.join(work_dir, "%s-cnv.bed" % (os.path.splitext(os.path.basename(base_bam))[0]))
    if not utils.file_exists(out_file):
        with file_transaction(out_file) as tx_out_file:
            with open(tx_out_file, "w") as out_handle:
                for chr_file in chr_files:
                    with open(chr_file) as in_handle:
                        is_empty = in_handle.readline().startswith("track name=empty")
                    if not is_empty:
                        with open(chr_file) as in_handle:
                            shutil.copyfileobj(in_handle, out_handle)
    return out_file

def _prep_sample_cnvs(cnv_file, data):
    """Convert a multiple sample CNV file into a single BED file for a sample.
    """
    import pybedtools
    sample_name = data["name"][-1]
    sample_file = os.path.join(os.path.dirname(cnv_file), "%s-cnv.bed" % sample_name)
    if not utils.file_exists(sample_file):
        with file_transaction(sample_file) as tx_out_file:
            pybedtools.BedTool(cnv_file).filter(lambda x: x.name == sample_name).saveas(tx_out_file)
    return sample_file

def _run_on_chrom(chrom, work_bams, names, work_dir, items):
    """Run cn.mops on work BAMs for a specific chromosome.
    """
    local_sitelib = os.path.join(install.get_defaults().get("tooldir", "/usr/local"),
                                 "lib", "R", "site-library")
    out_file = os.path.join(work_dir, "%s-%s-cnv.bed" % (os.path.splitext(os.path.basename(work_bams[0]))[0],
                                                         chrom))
    if not utils.file_exists(out_file):
        with file_transaction(out_file) as tx_out_file:
            rcode = "%s-run.R" % os.path.splitext(out_file)[0]
            with open(rcode, "w") as out_handle:
                out_handle.write(_script.format(prep_str=_prep_load_script(work_bams, names, chrom, items),
                                                out_file=tx_out_file,
                                                local_sitelib=local_sitelib))
            rscript = config_utils.get_program("Rscript", items[0]["config"])
            try:
                do.run([rscript, rcode], "cn.mops CNV detection", items[0], log_error=False)
            except subprocess.CalledProcessError, msg:
                # cn.mops errors out if no CNVs found. Just write an empty file.
                if _allowed_cnmops_errorstates(str(msg)):
                    with open(tx_out_file, "w") as out_handle:
                        out_handle.write('track name=empty description="No CNVs found"\n')
                else:
                    logger.exception()
                    raise
    return out_file

def _allowed_cnmops_errorstates(msg):
    return (msg.find("No CNV regions in result object. Rerun cn.mops with different parameters") >= 0
            or msg.find("Normalization might not be applicable for this small number of segments") >= 0
            or msg.find("Error in if (is.finite(mv2m)) { : argument is of length zero") >= 0)

def _prep_load_script(work_bams, names, chrom, items):
    pairmode = "paired" if bam.is_paired(work_bams[0]) else "unpaired"
    print len(items), items[0].get("metadata")
    if len(items) == 2 and vcfutils.get_paired_phenotype(items[0]):
        load_script = _paired_load_script
    else:
        load_script = _population_load_script
    return load_script(work_bams, names, chrom, pairmode, items)

def _population_load_script(work_bams, names, chrom, pairmode, items):
    """Prepare BAMs for assessing CNVs in a population.
    """
    return _population_prep.format(bam_file_str=",".join(work_bams), names_str=",".join(names),
                                   chrom=chrom, num_cores=0, pairmode=pairmode)

def _paired_load_script(work_bams, names, chrom, pairmode, items):
    """Prepare BAMs for assessing CNVs in a paired tumor/normal setup.
    """
    paired = vcfutils.get_paired_bams(work_bams, items)
    return _paired_prep.format(case_file=paired.tumor_bam, case_name=paired.tumor_name,
                               ctrl_file=paired.normal_bam, ctrl_name=paired.normal_name,
                               num_cores=0, chrom=chrom, pairmode=pairmode)

_script = """
.libPaths(c("{local_sitelib}"))
library(cn.mops)
library(rtracklayer)

{prep_str}

calc_cnvs <- cnvs(cnv_out)
strcn_to_cn <- function(x) {{
  as.numeric(substring(x, 3, 20))}}
calc_cnvs$score <- strcn_to_cn(calc_cnvs$CN)
calc_cnvs$name <- calc_cnvs$sampleName
export.bed(calc_cnvs, "{out_file}")
"""

_population_prep = """
bam_files <- strsplit("{bam_file_str}", ",")[[1]]
sample_names <- strsplit("{names_str}", ",")[[1]]
count_drs <- getReadCountsFromBAM(bam_files, sampleNames=sample_names, mode="{pairmode}",
                                  refSeqName="{chrom}", parallel={num_cores})
prep_counts <- cn.mops(count_drs, parallel={num_cores})
cnv_out <- calcIntegerCopyNumbers(prep_counts)
"""

_paired_prep = """
case_count <- getReadCountsFromBAM(c("{case_file}"), sampleNames=c("{case_name}"), mode="{pairmode}",
                                   refSeqName="{chrom}", parallel={num_cores})
ctrl_count <- getReadCountsFromBAM(c("{ctrl_file}"), sampleNames=c("{ctrl_name}"), mode="{pairmode}",
                                   refSeqName="{chrom}", parallel={num_cores},
                                   WL=width(case_count)[[1]])
prep_counts <- referencecn.mops(case_count, ctrl_count, parallel={num_cores})
cnv_out <- calcFractionalCopyNumbers(prep_counts)
"""
<commit_msg>Modified bcbio/structural/cn_mops.py to accept regions rather than always using the whole chromosomes.<commit_after>"""Copy number detection using read counts, with cn.mops.

http://www.bioconductor.org/packages/release/bioc/html/cn.mops.html
"""
from contextlib import closing
import os
import shutil
import subprocess

import pysam

from bcbio import bam, install, utils
from bcbio.distributed.transaction import file_transaction
from bcbio.log import logger
from bcbio.pipeline import config_utils
from bcbio.provenance import do
from bcbio.variation import vcfutils

def run(items):
    """Detect copy number variations from batched set of samples using cn.mops.
    """
    names = [x["name"][-1] for x in items]
    work_bams = [x["align_bam"] for x in items]
    if len(items) == 1:
        raise ValueError("cn.mops only works on batches with multiple samples")
    data = items[0]
    # XXX Can parallelize cn.mops with snow but currently causing
    # errors. Could also parallelize by chromosome.
    #num_cores = data["config"]["algorithm"].get("num_cores", 1)

    work_dir = utils.safe_makedir(os.path.join(data["dirs"]["work"], "structural", names[0],
                                               "cn_mops"))
    with closing(pysam.Samfile(work_bams[0], "rb")) as pysam_work_bam:
        out_files = [_run_on_chrom(chrom, work_bams, names, work_dir, items)
                     for chrom in pysam_work_bam.references]
    out_file = _combine_out_files(out_files, work_bams[0], work_dir)
    out = []
    for data in items:
        if "sv" not in data:
            data["sv"] = {}
        data["sv"]["cnv"] = _prep_sample_cnvs(out_file, data)
        out.append(data)
    return out

def _combine_out_files(chr_files, base_bam, work_dir):
    """Concatenate all CNV calls into a single file.
    """
    out_file = os.path.join(work_dir, "%s-cnv.bed" % (os.path.splitext(os.path.basename(base_bam))[0]))
    if not utils.file_exists(out_file):
        with file_transaction(out_file) as tx_out_file:
            with open(tx_out_file, "w") as out_handle:
                for chr_file in chr_files:
                    with open(chr_file) as in_handle:
                        is_empty = in_handle.readline().startswith("track name=empty")
                    if not is_empty:
                        with open(chr_file) as in_handle:
                            shutil.copyfileobj(in_handle, out_handle)
    return out_file

def _prep_sample_cnvs(cnv_file, data):
    """Convert a multiple sample CNV file into a single BED file for a sample.
    """
    import pybedtools
    sample_name = data["name"][-1]
    sample_file = os.path.join(os.path.dirname(cnv_file), "%s-cnv.bed" % sample_name)
    if not utils.file_exists(sample_file):
        with file_transaction(sample_file) as tx_out_file:
            pybedtools.BedTool(cnv_file).filter(lambda x: x.name == sample_name).saveas(tx_out_file)
    return sample_file

def _run_on_chrom(chrom, work_bams, names, work_dir, items):
    """Run cn.mops on work BAMs for a specific chromosome.
    """
    local_sitelib = os.path.join(install.get_defaults().get("tooldir", "/usr/local"),
                                 "lib", "R", "site-library")
    out_file = os.path.join(work_dir, "%s-%s-cnv.bed" % (os.path.splitext(os.path.basename(work_bams[0]))[0],
                                                         chrom))
    if not utils.file_exists(out_file):
        with file_transaction(out_file) as tx_out_file:
            rcode = "%s-run.R" % os.path.splitext(out_file)[0]
            with open(rcode, "w") as out_handle:
                out_handle.write(_script.format(prep_str=_prep_load_script(work_bams, names, chrom, items),
                                                out_file=tx_out_file,
                                                local_sitelib=local_sitelib))
            rscript = config_utils.get_program("Rscript", items[0]["config"])
            try:
                do.run([rscript, rcode], "cn.mops CNV detection", items[0], log_error=False)
            except subprocess.CalledProcessError, msg:
                # cn.mops errors out if no CNVs found. Just write an empty file.
                if _allowed_cnmops_errorstates(str(msg)):
                    with open(tx_out_file, "w") as out_handle:
                        out_handle.write('track name=empty description="No CNVs found"\n')
                else:
                    logger.exception()
                    raise
    return out_file

def _allowed_cnmops_errorstates(msg):
    return (msg.find("No CNV regions in result object. Rerun cn.mops with different parameters") >= 0
            or msg.find("Normalization might not be applicable for this small number of segments") >= 0
            or msg.find("Error in if (is.finite(mv2m)) { : argument is of length zero") >= 0)

def _prep_load_script(work_bams, names, chrom, items):
    pairmode = "paired" if bam.is_paired(work_bams[0]) else "unpaired"
    print len(items), items[0].get("metadata")
    if len(items) == 2 and vcfutils.get_paired_phenotype(items[0]):
        load_script = _paired_load_script
    else:
        load_script = _population_load_script
    return load_script(work_bams, names, chrom, pairmode, items)

def _population_load_script(work_bams, names, chrom, pairmode, items):
    """Prepare BAMs for assessing CNVs in a population.
    """
    bed_file = items[0]["config"]["algorithm"].get("variant_regions", None)
    if utils.file_exists(bed_file):
        return _population_prep_targeted.format(bam_file_str=",".join(work_bams), names_str=",".join(names),
                                                chrom=chrom, num_cores=0, pairmode=pairmode, bed_file=bed_file)
    else:
        return _population_prep.format(bam_file_str=",".join(work_bams), names_str=",".join(names),
                                       chrom=chrom, num_cores=0, pairmode=pairmode)

def _paired_load_script(work_bams, names, chrom, pairmode, items):
    """Prepare BAMs for assessing CNVs in a paired tumor/normal setup.
    """
    paired = vcfutils.get_paired_bams(work_bams, items)
    bed_file = items[0]["config"]["algorithm"].get("variant_regions", None)
    import pdb; pdb.set_trace()
    if utils.file_exists(bed_file):
        return _paired_prep_targeted.format(case_file=paired.tumor_bam, case_name=paired.tumor_name,
                                            ctrl_file=paired.normal_bam, ctrl_name=paired.normal_name,
                                            num_cores=0, chrom=chrom, pairmode=pairmode, bed_file=bed_file)
    else:
        return _paired_prep.format(case_file=paired.tumor_bam, case_name=paired.tumor_name,
                                   ctrl_file=paired.normal_bam, ctrl_name=paired.normal_name,
                                   num_cores=0, chrom=chrom, pairmode=pairmode)

_script = """
.libPaths(c("{local_sitelib}"))
library(cn.mops)
library(rtracklayer)

{prep_str}

calc_cnvs <- cnvs(cnv_out)
strcn_to_cn <- function(x) {{
  as.numeric(substring(x, 3, 20))}}
calc_cnvs$score <- strcn_to_cn(calc_cnvs$CN)
calc_cnvs$name <- calc_cnvs$sampleName
export.bed(calc_cnvs, "{out_file}")
"""

_population_prep = """
bam_files <- strsplit("{bam_file_str}", ",")[[1]]
sample_names <- strsplit("{names_str}", ",")[[1]]
count_drs <- getReadCountsFromBAM(bam_files, sampleNames=sample_names, mode="{pairmode}",
                                  refSeqName="{chrom}", parallel={num_cores})
prep_counts <- cn.mops(count_drs, parallel={num_cores})
cnv_out <- calcIntegerCopyNumbers(prep_counts)
"""

_paired_prep = """
case_count <- getReadCountsFromBAM(c("{case_file}"), sampleNames=c("{case_name}"), mode="{pairmode}",
                                   refSeqName="{chrom}", parallel={num_cores})
ctrl_count <- getReadCountsFromBAM(c("{ctrl_file}"), sampleNames=c("{ctrl_name}"), mode="{pairmode}",
                                   refSeqName="{chrom}", parallel={num_cores},
                                   WL=width(case_count)[[1]])
prep_counts <- referencecn.mops(case_count, ctrl_count, parallel={num_cores})
cnv_out <- calcFractionalCopyNumbers(prep_counts)
"""

_population_prep_targeted = """
bam_files <- strsplit("{bam_file_str}", ",")[[1]]
sample_names <- strsplit("{names_str}", ",")[[1]]
my_gr <- import.bed(c("{bed_file}"), trackLine=FALSE)
my_gr_subset = subset(my_gr, seqnames(my_gr) == "{chrom}")
if (length(my_gr_subset) < 1) stop("No CNV regions in result object. Rerun cn.mops with different parameters!")
count_drs <- getSegmentReadCountsFromBAM(bam_files, sampleNames=sample_names, mode="{pairmode}",
                                         GR=my_gr, parallel={num_cores})
prep_counts <- cn.mops(count_drs, parallel={num_cores})
cnv_out <- calcIntegerCopyNumbers(prep_counts)
"""

_paired_prep_targeted = """
my_gr <- import.bed(c("{bed_file}"), trackLine=FALSE)
my_gr_subset = subset(my_gr, seqnames(my_gr) == "{chrom}")
if (length(my_gr_subset) < 1) stop("No CNV regions in result object. Rerun cn.mops with different parameters!")
case_count <- getSegmentReadCountsFromBAM(c("{case_file}"), GR=my_gr,
                                          sampleNames=c("{case_name}"),
                                          mode="{pairmode}", parallel={num_cores})
ctrl_count <- getSegmentReadCountsFromBAM(c("{ctrl_file}"), GR=my_gr,
                                          sampleNames=c("{case_name}"),
                                          mode="{pairmode}", parallel={num_cores})
prep_counts <- referencecn.mops(case_count, ctrl_count, parallel={num_cores})
cnv_out <- calcFractionalCopyNumbers(prep_counts)
"""
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright (C) 2006-2007 Edgewall Software
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://genshi.edgewall.org/wiki/License.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://genshi.edgewall.org/log/.

"""Implementation of a number of stream filters."""

try:
    frozenset
except NameError:
    from sets import ImmutableSet as frozenset
import re

from genshi.core import Attrs, QName, stripentities
from genshi.core import END, START, TEXT

__all__ = ['HTMLFormFiller', 'HTMLSanitizer']

class HTMLFormFiller(object):
    """A stream filter that can populate HTML forms from a dictionary of values.

    >>> from genshi.input import HTML
    >>> html = HTML('''<form>
    ...
<input type="text" name="foo" />

    ... </form>''')
    >>> filler = HTMLFormFiller(data={'foo': 'bar'})
    >>> print html | filler
    <form>

<input type="text" name="foo" value="bar"/>

    </form>
    """
    # TODO: only select the first radio button, and the first select option
    #       (if not in a multiple-select)
    # TODO: only apply to elements in the XHTML namespace (or no namespace)?

    def __init__(self, name=None, id=None, data=None):
        """Create the filter.

        @param name: The name of the form that should be populated. If this
            parameter is given, only forms where the ``name`` attribute value
            matches the parameter are processed.
        @param id: The ID of the form that should be populated. If this
            parameter is given, only forms where the ``id`` attribute value
            matches the parameter are processed.
        @param data: The dictionary of form values, where the keys are the names
            of the form fields, and the values are the values to fill in.
        """
        self.name = name
        self.id = id
        if data is None:
            data = {}
        self.data = data

    def __call__(self, stream, ctxt=None):
        """Apply the filter to the given stream.

        @param stream: the markup event stream to filter
        @param ctxt: the template context (unused)
        """
        in_form = in_select = in_option = in_textarea = False
        select_value = option_value = textarea_value = None
        option_start = option_text = None

        for kind, data, pos in stream:

            if kind is START:
                tag, attrs = data
                tagname = tag.localname

                if tagname == 'form' and (
                        self.name and attrs.get('name') == self.name or
                        self.id and attrs.get('id') == self.id or
                        not (self.id or self.name)):
                    in_form = True

                elif in_form:
                    if tagname == 'input':
                        type = attrs.get('type')
                        if type in ('checkbox', 'radio'):
                            name = attrs.get('name')
                            if name:
                                value = self.data.get(name)
                                declval = attrs.get('value')
                                checked = False
                                if isinstance(value, (list, tuple)):
                                    if declval:
                                        checked = declval in value
                                    else:
                                        checked = bool(filter(None, value))
                                else:
                                    if declval:
                                        checked = declval == value
                                    elif type == 'checkbox':
                                        checked = bool(value)
                                if checked:
                                    attrs |= [(QName('checked'), 'checked')]
                                elif 'checked' in attrs:
                                    attrs -= 'checked'
                        elif type in (None, 'hidden', 'text'):
                            name = attrs.get('name')
                            if name:
                                value = self.data.get(name)
                                if isinstance(value, (list, tuple)):
                                    value = value[0]
                                if value is not None:
                                    attrs |= [(QName('value'), unicode(value))]
                    elif tagname == 'select':
                        name = attrs.get('name')
                        select_value = self.data.get(name)
                        in_select = True
                    elif tagname == 'textarea':
                        name = attrs.get('name')
                        textarea_value = self.data.get(name)
                        if isinstance(textarea_value, (list, tuple)):
                            textarea_value = textarea_value[0]
                        in_textarea = True
                    elif in_select and tagname == 'option':
                        option_start = kind, data, pos
                        option_value = attrs.get('value')
                        in_option = True
                        continue
                yield kind, (tag, attrs), pos

            elif in_form and kind is TEXT:
                if in_select and in_option:
                    if option_value is None:
                        option_value = data
                    option_text = kind, data, pos
                    continue
                elif in_textarea:
                    continue
                yield kind, data, pos

            elif in_form and kind is END:
                tagname = data.localname
                if tagname == 'form':
                    in_form = False
                elif tagname == 'select':
                    in_select = False
                    select_value = None
                elif in_select and tagname == 'option':
                    if isinstance(select_value, (tuple, list)):
                        selected = option_value in select_value
                    else:
                        selected = option_value == select_value
                    okind, (tag, attrs), opos = option_start
                    if selected:
                        attrs |= [(QName('selected'), 'selected')]
                    elif 'selected' in attrs:
                        attrs -= 'selected'
                    yield okind, (tag, attrs), opos
                    if option_text:
                        yield option_text
                    in_option = False
                    option_start = option_text = option_value = None
                elif tagname == 'textarea':
                    if textarea_value:
                        yield TEXT, unicode(textarea_value), pos
                    in_textarea = False
                yield kind, data, pos

            else:
                yield kind, data, pos

class HTMLSanitizer(object):
    """A filter that removes potentially dangerous HTML tags and attributes
    from the stream.
    """

    SAFE_TAGS = frozenset(['a', 'abbr', 'acronym', 'address', 'area', 'b',
        'big', 'blockquote', 'br', 'button', 'caption', 'center', 'cite',
        'code', 'col', 'colgroup', 'dd', 'del', 'dfn', 'dir', 'div', 'dl', 'dt',
        'em', 'fieldset', 'font', 'form', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6',
        'hr', 'i', 'img', 'input', 'ins', 'kbd', 'label', 'legend', 'li', 'map',
        'menu', 'ol', 'optgroup', 'option', 'p', 'pre', 'q', 's', 'samp',
        'select', 'small', 'span', 'strike', 'strong', 'sub', 'sup', 'table',
        'tbody', 'td', 'textarea', 'tfoot', 'th', 'thead', 'tr', 'tt', 'u',
        'ul', 'var'])

    SAFE_ATTRS = frozenset(['abbr', 'accept', 'accept-charset', 'accesskey',
        'action', 'align', 'alt', 'axis', 'bgcolor', 'border', 'cellpadding',
        'cellspacing', 'char', 'charoff', 'charset', 'checked', 'cite', 'class',
        'clear', 'cols', 'colspan', 'color', 'compact', 'coords', 'datetime',
        'dir', 'disabled', 'enctype', 'for', 'frame', 'headers', 'height',
        'href', 'hreflang', 'hspace', 'id', 'ismap', 'label', 'lang',
        'longdesc', 'maxlength', 'media', 'method', 'multiple', 'name',
        'nohref', 'noshade', 'nowrap', 'prompt', 'readonly', 'rel', 'rev',
        'rows', 'rowspan', 'rules', 'scope', 'selected', 'shape', 'size',
        'span', 'src', 'start', 'style', 'summary', 'tabindex', 'target',
        'title', 'type', 'usemap', 'valign', 'value', 'vspace', 'width'])

    SAFE_SCHEMES = frozenset(['file', 'ftp', 'http', 'https', 'mailto', None])

    URI_ATTRS = frozenset(['action', 'background', 'dynsrc', 'href', 'lowsrc',
        'src'])

    def __init__(self, safe_tags=SAFE_TAGS, safe_attrs=SAFE_ATTRS,
                 safe_schemes=SAFE_SCHEMES, uri_attrs=URI_ATTRS):
        """Create the sanitizer.

        The exact set of allowed elements and attributes can be configured.

        @param safe_tags: a set of tag names that are considered safe
        @param safe_attrs: a set of attribute names that are considered safe
        @param safe_schemes: a set of URI schemes that are considered safe
        @param uri_attrs: a set of names of attributes that contain URIs
        """
        self.safe_tags = safe_tags
        self.safe_attrs = safe_attrs
        self.uri_attrs = uri_attrs
        self.safe_schemes = safe_schemes

    def __call__(self, stream, ctxt=None):
        """Apply the filter to the given stream.

        @param stream: the markup event stream to filter
        @param ctxt: the template context (unused)
        """
        waiting_for = None

        def _get_scheme(href):
            if ':' not in href:
                return None
            chars = [char for char in href.split(':', 1)[0] if char.isalnum()]
            return ''.join(chars).lower()

        for kind, data, pos in stream:
            if kind is START:
                if waiting_for:
                    continue
                tag, attrs = data
                if tag not in self.safe_tags:
                    waiting_for = tag
                    continue

                new_attrs = []
                for attr, value in attrs:
                    value = stripentities(value)
                    if attr not in self.safe_attrs:
                        continue
                    elif attr in self.uri_attrs:
                        # Don't allow URI schemes such as "javascript:"
                        if _get_scheme(value) not in self.safe_schemes:
                            continue
                    elif attr == 'style':
                        # Remove dangerous CSS declarations from inline styles
                        decls = []
                        for decl in filter(None, value.split(';')):
                            is_evil = False
                            if 'expression' in decl:
                                is_evil = True
                            for m in re.finditer(r'url\s*\(([^)]+)', decl):
                                if _get_scheme(m.group(1)) not in self.safe_schemes:
                                    is_evil = True
                                    break
                            if not is_evil:
                                decls.append(decl.strip())
                        if not decls:
                            continue
                        value = '; '.join(decls)
                    new_attrs.append((attr, value))

                yield kind, (tag, Attrs(new_attrs)), pos

            elif kind is END:
                tag = data
                if waiting_for:
                    if waiting_for == tag:
                        waiting_for = None
                else:
                    yield kind, data, pos

            else:
                if not waiting_for:
                    yield kind, data, pos
<commit_msg>`HTMLFormFiller` now correctly deals with non-string values in the data dictionary for select/checkbox/radio controls.<commit_after># -*- coding: utf-8 -*-
#
# Copyright (C) 2006-2007 Edgewall Software
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://genshi.edgewall.org/wiki/License.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://genshi.edgewall.org/log/.

"""Implementation of a number of stream filters."""

try:
    frozenset
except NameError:
    from sets import ImmutableSet as frozenset
import re

from genshi.core import Attrs, QName, stripentities
from genshi.core import END, START, TEXT

__all__ = ['HTMLFormFiller', 'HTMLSanitizer']

class HTMLFormFiller(object):
    """A stream filter that can populate HTML forms from a dictionary of values.

    >>> from genshi.input import HTML
    >>> html = HTML('''<form>
    ...
<input type="text" name="foo" />

    ... </form>''')
    >>> filler = HTMLFormFiller(data={'foo': 'bar'})
    >>> print html | filler
    <form>

<input type="text" name="foo" value="bar"/>

    </form>
    """
    # TODO: only select the first radio button, and the first select option
    #       (if not in a multiple-select)
    # TODO: only apply to elements in the XHTML namespace (or no namespace)?

    def __init__(self, name=None, id=None, data=None):
        """Create the filter.

        @param name: The name of the form that should be populated. If this
            parameter is given, only forms where the ``name`` attribute value
            matches the parameter are processed.
        @param id: The ID of the form that should be populated. If this
            parameter is given, only forms where the ``id`` attribute value
            matches the parameter are processed.
        @param data: The dictionary of form values, where the keys are the names
            of the form fields, and the values are the values to fill in.
        """
        self.name = name
        self.id = id
        if data is None:
            data = {}
        self.data = data

    def __call__(self, stream, ctxt=None):
        """Apply the filter to the given stream.

        @param stream: the markup event stream to filter
        @param ctxt: the template context (unused)
        """
        in_form = in_select = in_option = in_textarea = False
        select_value = option_value = textarea_value = None
        option_start = option_text = None

        for kind, data, pos in stream:

            if kind is START:
                tag, attrs = data
                tagname = tag.localname

                if tagname == 'form' and (
                        self.name and attrs.get('name') == self.name or
                        self.id and attrs.get('id') == self.id or
                        not (self.id or self.name)):
                    in_form = True

                elif in_form:
                    if tagname == 'input':
                        type = attrs.get('type')
                        if type in ('checkbox', 'radio'):
                            name = attrs.get('name')
                            if name:
                                value = self.data.get(name)
                                declval = attrs.get('value')
                                checked = False
                                if isinstance(value, (list, tuple)):
                                    if declval:
                                        checked = declval in [str(v) for v
                                                              in value]
                                    else:
                                        checked = bool(filter(None, value))
                                else:
                                    if declval:
                                        checked = declval == str(value)
                                    elif type == 'checkbox':
                                        checked = bool(value)
                                if checked:
                                    attrs |= [(QName('checked'), 'checked')]
                                elif 'checked' in attrs:
                                    attrs -= 'checked'
                        elif type in (None, 'hidden', 'text'):
                            name = attrs.get('name')
                            if name:
                                value = self.data.get(name)
                                if isinstance(value, (list, tuple)):
                                    value = value[0]
                                if value is not None:
                                    attrs |= [(QName('value'), unicode(value))]
                    elif tagname == 'select':
                        name = attrs.get('name')
                        select_value = self.data.get(name)
                        in_select = True
                    elif tagname == 'textarea':
                        name = attrs.get('name')
                        textarea_value = self.data.get(name)
                        if isinstance(textarea_value, (list, tuple)):
                            textarea_value = textarea_value[0]
                        in_textarea = True
                    elif in_select and tagname == 'option':
                        option_start = kind, data, pos
                        option_value = attrs.get('value')
                        in_option = True
                        continue
                yield kind, (tag, attrs), pos

            elif in_form and kind is TEXT:
                if in_select and in_option:
                    if option_value is None:
                        option_value = data
                    option_text = kind, data, pos
                    continue
                elif in_textarea:
                    continue
                yield kind, data, pos

            elif in_form and kind is END:
                tagname = data.localname
                if tagname == 'form':
                    in_form = False
                elif tagname == 'select':
                    in_select = False
                    select_value = None
                elif in_select and tagname == 'option':
                    if isinstance(select_value, (tuple, list)):
                        selected = option_value in [str(v) for v
                                                    in select_value]
                    else:
                        selected = option_value == str(select_value)
                    okind, (tag, attrs), opos = option_start
                    if selected:
                        attrs |= [(QName('selected'), 'selected')]
                    elif 'selected' in attrs:
                        attrs -= 'selected'
                    yield okind, (tag, attrs), opos
                    if option_text:
                        yield option_text
                    in_option = False
                    option_start = option_text = option_value = None
                elif tagname == 'textarea':
                    if textarea_value:
                        yield TEXT, unicode(textarea_value), pos
                    in_textarea = False
                yield kind, data, pos

            else:
                yield kind, data, pos

class HTMLSanitizer(object):
    """A filter that removes potentially dangerous HTML tags and attributes
    from the stream.
    """

    SAFE_TAGS = frozenset(['a', 'abbr', 'acronym', 'address', 'area', 'b',
        'big', 'blockquote', 'br', 'button', 'caption', 'center', 'cite',
        'code', 'col', 'colgroup', 'dd', 'del', 'dfn', 'dir', 'div', 'dl', 'dt',
        'em', 'fieldset', 'font', 'form', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6',
        'hr', 'i', 'img', 'input', 'ins', 'kbd', 'label', 'legend', 'li', 'map',
        'menu', 'ol', 'optgroup', 'option', 'p', 'pre', 'q', 's', 'samp',
        'select', 'small', 'span', 'strike', 'strong', 'sub', 'sup', 'table',
        'tbody', 'td', 'textarea', 'tfoot', 'th', 'thead', 'tr', 'tt', 'u',
        'ul', 'var'])

    SAFE_ATTRS = frozenset(['abbr', 'accept', 'accept-charset', 'accesskey',
        'action', 'align', 'alt', 'axis', 'bgcolor', 'border', 'cellpadding',
        'cellspacing', 'char', 'charoff', 'charset', 'checked', 'cite', 'class',
        'clear', 'cols', 'colspan', 'color', 'compact', 'coords', 'datetime',
        'dir', 'disabled', 'enctype', 'for', 'frame', 'headers', 'height',
        'href', 'hreflang', 'hspace', 'id', 'ismap', 'label', 'lang',
        'longdesc', 'maxlength', 'media', 'method', 'multiple', 'name',
        'nohref', 'noshade', 'nowrap', 'prompt', 'readonly', 'rel', 'rev',
        'rows', 'rowspan', 'rules', 'scope', 'selected', 'shape', 'size',
        'span', 'src', 'start', 'style', 'summary', 'tabindex', 'target',
        'title', 'type', 'usemap', 'valign', 'value', 'vspace', 'width'])

    SAFE_SCHEMES = frozenset(['file', 'ftp', 'http', 'https', 'mailto', None])

    URI_ATTRS = frozenset(['action', 'background', 'dynsrc', 'href', 'lowsrc',
        'src'])

    def __init__(self, safe_tags=SAFE_TAGS, safe_attrs=SAFE_ATTRS,
                 safe_schemes=SAFE_SCHEMES, uri_attrs=URI_ATTRS):
        """Create the sanitizer.

        The exact set of allowed elements and attributes can be configured.

        @param safe_tags: a set of tag names that are considered safe
        @param safe_attrs: a set of attribute names that are considered safe
        @param safe_schemes: a set of URI schemes that are considered safe
        @param uri_attrs: a set of names of attributes that contain URIs
        """
        self.safe_tags = safe_tags
        self.safe_attrs = safe_attrs
        self.uri_attrs = uri_attrs
        self.safe_schemes = safe_schemes

    def __call__(self, stream, ctxt=None):
        """Apply the filter to the given stream.

        @param stream: the markup event stream to filter
        @param ctxt: the template context (unused)
        """
        waiting_for = None

        def _get_scheme(href):
            if ':' not in href:
                return None
            chars = [char for char in href.split(':', 1)[0] if char.isalnum()]
            return ''.join(chars).lower()

        for kind, data, pos in stream:
            if kind is START:
                if waiting_for:
                    continue
                tag, attrs = data
                if tag not in self.safe_tags:
                    waiting_for = tag
                    continue

                new_attrs = []
                for attr, value in attrs:
                    value = stripentities(value)
                    if attr not in self.safe_attrs:
                        continue
                    elif attr in self.uri_attrs:
                        # Don't allow URI schemes such as "javascript:"
                        if _get_scheme(value) not in self.safe_schemes:
                            continue
                    elif attr == 'style':
                        # Remove dangerous CSS declarations from inline styles
                        decls = []
                        for decl in filter(None, value.split(';')):
                            is_evil = False
                            if 'expression' in decl:
                                is_evil = True
                            for m in re.finditer(r'url\s*\(([^)]+)', decl):
                                if _get_scheme(m.group(1)) not in self.safe_schemes:
                                    is_evil = True
                                    break
                            if not is_evil:
                                decls.append(decl.strip())
                        if not decls:
                            continue
                        value = '; '.join(decls)
                    new_attrs.append((attr, value))

                yield kind, (tag, Attrs(new_attrs)), pos

            elif kind is END:
                tag = data
                if waiting_for:
                    if waiting_for == tag:
                        waiting_for = None
                else:
                    yield kind, data, pos

            else:
                if not waiting_for:
                    yield kind, data, pos
<|endoftext|>


<commit_before>__version_info__ = (1, 5, 0)
__version__ = '.'.join(str(v) for v in __version_info__)
__git_revno__ = "b350b9257de3481d7787ba8993a672245dd29549"
<commit_msg>v1.5.1<commit_after>__version_info__ = (1, 5, 1)
__version__ = '.'.join(str(v) for v in __version_info__)
__git_revno__ = "b350b9257de3481d7787ba8993a672245dd29549"
<|endoftext|>