1#pylint: disable=attribute-defined-outside-init
2from __future__ import division
3import csv
4import os
5import time
6import tempfile
7from fcntl import fcntl, F_GETFL, F_SETFL
8from string import Template
9from subprocess import Popen, PIPE, STDOUT
10
11from devlib import Instrument, CONTINUOUS, MeasurementsCsv
12from devlib.exception import HostError
13from devlib.utils.misc import which
14
15OUTPUT_CAPTURE_FILE = 'acme-cape.csv'
16IIOCAP_CMD_TEMPLATE = Template("""
17${iio_capture} -n ${host} -b ${buffer_size} -c -f ${outfile} ${iio_device}
18""")
19
20def _read_nonblock(pipe, size=1024):
21    fd = pipe.fileno()
22    flags = fcntl(fd, F_GETFL)
23    flags |= os.O_NONBLOCK
24    fcntl(fd, F_SETFL, flags)
25
26    output = ''
27    try:
28        while True:
29            output += pipe.read(size)
30    except IOError:
31        pass
32    return output
33
34
35class AcmeCapeInstrument(Instrument):
36
37    mode = CONTINUOUS
38
39    def __init__(self, target,
40                 iio_capture=which('iio-capture'),
41                 host='baylibre-acme.local',
42                 iio_device='iio:device0',
43                 buffer_size=256):
44        super(AcmeCapeInstrument, self).__init__(target)
45        self.iio_capture = iio_capture
46        self.host = host
47        self.iio_device = iio_device
48        self.buffer_size = buffer_size
49        self.sample_rate_hz = 100
50        if self.iio_capture is None:
51            raise HostError('Missing iio-capture binary')
52        self.command = None
53        self.process = None
54
55        self.add_channel('shunt', 'voltage')
56        self.add_channel('bus', 'voltage')
57        self.add_channel('device', 'power')
58        self.add_channel('device', 'current')
59        self.add_channel('timestamp', 'time_ms')
60
61    def __del__(self):
62        if self.process and self.process.pid:
63            self.logger.warning('killing iio-capture process [%d]...',
64                                self.process.pid)
65            self.process.kill()
66
67    def reset(self, sites=None, kinds=None, channels=None):
68        super(AcmeCapeInstrument, self).reset(sites, kinds, channels)
69        self.raw_data_file = tempfile.mkstemp('.csv')[1]
70        params = dict(
71            iio_capture=self.iio_capture,
72            host=self.host,
73            buffer_size=self.buffer_size,
74            iio_device=self.iio_device,
75            outfile=self.raw_data_file
76        )
77        self.command = IIOCAP_CMD_TEMPLATE.substitute(**params)
78        self.logger.debug('ACME cape command: {}'.format(self.command))
79
80    def start(self):
81        self.process = Popen(self.command.split(), stdout=PIPE, stderr=STDOUT)
82
83    def stop(self):
84        self.process.terminate()
85        timeout_secs = 10
86        output = ''
87        for _ in xrange(timeout_secs):
88            if self.process.poll() is not None:
89                break
90            time.sleep(1)
91        else:
92            output += _read_nonblock(self.process.stdout)
93            self.process.kill()
94            self.logger.error('iio-capture did not terminate gracefully')
95            if self.process.poll() is None:
96                msg = 'Could not terminate iio-capture:\n{}'
97                raise HostError(msg.format(output))
98        if self.process.returncode != 15: # iio-capture exits with 15 when killed
99            output += self.process.stdout.read()
100            self.logger.info('ACME instrument encountered an error, '
101                             'you may want to try rebooting the ACME device:\n'
102                             '  ssh root@{} reboot'.format(self.host))
103            raise HostError('iio-capture exited with an error ({}), output:\n{}'
104                            .format(self.process.returncode, output))
105        if not os.path.isfile(self.raw_data_file):
106            raise HostError('Output CSV not generated.')
107        self.process = None
108
109    def get_data(self, outfile):
110        if os.stat(self.raw_data_file).st_size == 0:
111            self.logger.warning('"{}" appears to be empty'.format(self.raw_data_file))
112            return
113
114        all_channels = [c.label for c in self.list_channels()]
115        active_channels = [c.label for c in self.active_channels]
116        active_indexes = [all_channels.index(ac) for ac in active_channels]
117
118        with open(self.raw_data_file, 'rb') as fh:
119            with open(outfile, 'wb') as wfh:
120                writer = csv.writer(wfh)
121                writer.writerow(active_channels)
122
123                reader = csv.reader(fh, skipinitialspace=True)
124                header = reader.next()
125                ts_index = header.index('timestamp ms')
126
127
128                for row in reader:
129                    output_row = []
130                    for i in active_indexes:
131                        if i == ts_index:
132                            # Leave time in ms
133                            output_row.append(float(row[i]))
134                        else:
135                            # Convert rest into standard units.
136                            output_row.append(float(row[i])/1000)
137                    writer.writerow(output_row)
138        return MeasurementsCsv(outfile, self.active_channels, self.sample_rate_hz)
139
140    def get_raw(self):
141        return [self.raw_data_file]
142