1#!/usr/bin/env python
2
3from grammar import resd_header, data_block, data_block_sample_frequency, data_block_sample_arbitrary, data_block_header, data_block_subheader, data_block_metadata_item, BLOCK_TYPE, SAMPLE_TYPE
4
5__VERSION__ = 1
6
7
8class RESD:
9    def __init__(self, file_path):
10        self.file_handle = open(file_path, 'wb')
11        self.blocks = {}
12        self._write_header()
13
14    def __del__(self):
15        self.flush()
16        self.file_handle.close()
17
18    def new_block(self, sample_type, block_type, channel_id=0):
19        previous_block = self.get_block(sample_type, channel_id)
20        if previous_block is not None:
21            self.flush(sample_type, channel_id)
22
23        block = ({
24            BLOCK_TYPE.CONSTANT_FREQUENCY: RESDBlockConstantFrequency,
25            BLOCK_TYPE.ARBITRARY_TIMESTAMP: RESDBlockArbitraryTimestamp
26        })[block_type](sample_type, block_type, channel_id)
27
28        self.blocks[(sample_type, channel_id)] = block
29        return block
30
31    def get_block(self, sample_type, channel_id=0):
32        return self.blocks.get((sample_type, channel_id), None)
33
34    def get_block_or_create(self, sample_type, block_type, channel_id=0):
35        block = self.get_block(sample_type, channel_id)
36        return block if block else self.new_block(sample_type, block_type, channel_id)
37
38    def flush(self, sample_type=None, channel_id=None):
39        for key in list(self.blocks.keys()):
40            block_sample_type, block_channel_id = key
41
42            if sample_type and block_sample_type != sample_type:
43                continue
44            if channel_id and block_channel_id != channel_id:
45                continue
46
47            self.blocks[key].flush(self.file_handle)
48            del self.blocks[key]
49
50    def _write_header(self):
51        resd_header.build_stream({
52            'version': __VERSION__,
53        }, self.file_handle)
54
55
56class RESDBlock:
57    def __init__(self, sample_type, block_type, channel_id):
58        self.sample_type = sample_type
59        self.block_type = block_type
60        self.channel_id = channel_id
61        self.block_metadata = RESDBlockMetadata()
62        self.samples = []
63
64    @property
65    def metadata(self):
66        return self.block_metadata
67
68    def flush(self, file):
69        metadata = self.metadata.build()
70        data_size = (
71            data_block_subheader.sizeof(header={'block_type': self.block_type}) +
72            metadata['size'] + 8 +
73            self._samples_sizeof()
74        )
75
76        header = self._header(data_size)
77        subheader = self._subheader()
78        data_block.build_stream({
79            'header': header,
80            'subheader': subheader,
81            'metadata': metadata,
82            'samples': self.samples,
83        }, file)
84
85    def _header(self, data_size):
86        return {
87            'block_type': self.block_type,
88            'sample_type': self.sample_type,
89            'channel_id': self.channel_id,
90            'data_size': data_size,
91        }
92
93    def _subheader(self):
94        return None
95
96    def _samples_sizeof(self):
97        pass
98
99
100class RESDBlockConstantFrequency(RESDBlock):
101    __period = int(1e9)
102    __start_time = 0
103
104    @property
105    def period(self):
106        return self.__period
107
108    @period.setter
109    def period(self, value):
110        self.__period = value
111
112    @property
113    def frequency(self):
114        return 1e9 / self.__period
115
116    @frequency.setter
117    def frequency(self, value):
118        self.__period = int(1e9 / value)
119
120    @property
121    def start_time(self):
122        return self.__start_time
123
124    @start_time.setter
125    def start_time(self, value):
126        self.__start_time = value
127
128    def add_sample(self, sample):
129        self.samples.append({'sample': sample})
130
131    def _subheader(self):
132        return {
133            'start_time': self.__start_time,
134            'period': self.__period
135        }
136
137    def _samples_sizeof(self):
138        return sum(len(data_block_sample_frequency(self.sample_type).build(sample)) for sample in self.samples)
139
140
141class RESDBlockArbitraryTimestamp(RESDBlock):
142    __start_time = 0
143
144    @property
145    def start_time(self):
146        return self.__start_time
147
148    @start_time.setter
149    def start_time(self, value):
150        self.__start_time = value
151
152    def add_sample(self, sample, timestamp):
153        self.samples.append({'sample': sample, 'timestamp': timestamp})
154
155    def _subheader(self):
156        return {
157            'start_time': self.__start_time,
158        }
159
160    def _samples_sizeof(self):
161        return sum(len(data_block_sample_arbitrary(self.sample_type).build(sample)) for sample in self.samples)
162
163
164class RESDBlockMetadata:
165    def __init__(self):
166        self.metadata = []
167        self.keys = set()
168
169    def __getattr__(self, name):
170        prefix = 'insert_'
171        if name[:len(prefix)] != prefix:
172            return None
173
174        method = name[len(prefix):]
175        type_idx = ({
176            'int8':   0x00,
177            'uint8':  0x01,
178            'int16':  0x02,
179            'uint16': 0x03,
180            'int32':  0x04,
181            'uint32': 0x05,
182            'int64':  0x06,
183            'uint64': 0x07,
184            'float':  0x08,
185            'double': 0x09,
186            'text':   0x0A,
187            'blob':   0x0B,
188        }).get(method, None)
189
190        if method is None:
191            return None
192
193        return lambda key, value: self._insert(type_idx, key, value)
194
195    def build(self):
196        return {'items': self.metadata, 'size': self._sizeof()}
197
198    def remove(self, key):
199        if key not in self.keys:
200            return
201        self.keys.remove(key)
202        index = next(i for i, value in enumerate(self.metadata) if value['key'] == key)
203        self.metadata.pop(index)
204
205    def _sizeof(self):
206        return sum(len(data_block_metadata_item.build(item)) for item in self.metadata)
207
208    def _insert(self, type_idx, key, value):
209        self.remove(key)
210        self.keys.add(key)
211        self.metadata.push({
212            'type': type_idx,
213            'key': key,
214            'value': value
215        })
216