cdcdd8956a06c860e2491b508027b2295428ecf0
[igor.git] / igor / record / variables.py
1 # Copyright (C) 2012 W. Trevor King <wking@tremily.us>
2 #
3 # This file is part of igor.
4 #
5 # igor is free software: you can redistribute it and/or modify it under the
6 # terms of the GNU Lesser General Public License as published by the Free
7 # Software Foundation, either version 3 of the License, or (at your option) any
8 # later version.
9 #
10 # igor is distributed in the hope that it will be useful, but WITHOUT ANY
11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
12 # A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
13 # details.
14 #
15 # You should have received a copy of the GNU Lesser General Public License
16 # along with igor.  If not, see <http://www.gnu.org/licenses/>.
17
18 import io as _io
19
20 from .. import LOG as _LOG
21 from ..binarywave import TYPE_TABLE as _TYPE_TABLE
22 from ..binarywave import NullStaticStringField as _NullStaticStringField
23 from ..binarywave import DynamicStringField as _DynamicStringField
24 from ..struct import Structure as _Structure
25 from ..struct import DynamicStructure as _DynamicStructure
26 from ..struct import Field as _Field
27 from ..struct import DynamicField as _DynamicField
28 from ..util import byte_order as _byte_order
29 from ..util import need_to_reorder_bytes as _need_to_reorder_bytes
30 from .base import Record
31
32
33 class ListedStaticStringField (_NullStaticStringField):
34     """Handle string conversions for multi-count dynamic parents.
35
36     If a field belongs to a multi-count dynamic parent, the parent is
37     called multiple times to parse each count, and the field's
38     post-unpack hook gets called after the field is unpacked during
39     each iteration.  This requires alternative logic for getting and
40     setting the string data.  The actual string formatting code is not
41     affected.
42     """
43     def post_unpack(self, parents, data):
44         parent_structure = parents[-1]
45         parent_data = self._get_structure_data(parents, data, parent_structure)
46         d = self._normalize_string(parent_data[-1][self.name])
47         parent_data[-1][self.name] = d
48
49
50 class ListedStaticStringField (_NullStaticStringField):
51     """Handle string conversions for multi-count dynamic parents.
52
53     If a field belongs to a multi-count dynamic parent, the parent is
54     called multiple times to parse each count, and the field's
55     post-unpack hook gets called after the field is unpacked during
56     each iteration.  This requires alternative logic for getting and
57     setting the string data.  The actual string formatting code is not
58     affected.
59     """
60     def post_unpack(self, parents, data):
61         parent_structure = parents[-1]
62         parent_data = self._get_structure_data(parents, data, parent_structure)
63         d = self._normalize_string(parent_data[-1][self.name])
64         parent_data[-1][self.name] = d
65
66
67 class ListedDynamicStrDataField (_DynamicStringField, ListedStaticStringField):
68     _size_field = 'strLen'
69     _null_terminated = False
70
71     def _get_size_data(self, parents, data):
72         parent_structure = parents[-1]
73         parent_data = self._get_structure_data(parents, data, parent_structure)
74         return parent_data[-1][self._size_field]
75
76
77 class DynamicVarDataField (_DynamicField):
78     def pre_pack(self, parents, data):
79         raise NotImplementedError()
80
81     def post_unpack(self, parents, data):
82         var_structure = parents[-1]
83         var_data = self._get_structure_data(parents, data, var_structure)
84         data = var_data[self.name]
85         d = {}
86         for i,value in enumerate(data):
87             key,value = self._normalize_item(i, value)
88             d[key] = value
89         var_data[self.name] = d
90
91     def _normalize_item(self, index, value):
92         raise NotImplementedError()
93
94
95 class DynamicSysVarField (DynamicVarDataField):
96     def _normalize_item(self, index, value):
97         name = 'K{}'.format(index)
98         return (name, value)
99
100
101 class DynamicUserVarField (DynamicVarDataField):
102     def _normalize_item(self, index, value):
103         name = value['name']
104         value = value['num']
105         return (name, value)
106
107
108 class DynamicUserStrField (DynamicVarDataField):
109     def _normalize_item(self, index, value):
110         name = value['name']
111         value = value['data']
112         return (name, value)
113
114
115 class DynamicVarNumField (_DynamicField):
116     def post_unpack(self, parents, data):
117         parent_structure = parents[-1]
118         parent_data = self._get_structure_data(parents, data, parent_structure)
119         d = self._normalize_numeric_variable(parent_data[-1][self.name])
120         parent_data[-1][self.name] = d
121
122     def _normalize_numeric_variable(self, num_var):
123         t = _TYPE_TABLE[num_var['numType']]
124         if num_var['numType'] % 2:  # complex number
125             return t(complex(num_var['realPart'], num_var['imagPart']))
126         else:
127             return t(num_var['realPart'])
128
129
130 class DynamicFormulaField (_DynamicStringField):
131     _size_field = 'formulaLen'
132     _null_terminated = True
133
134
135 # From Variables.h
136 VarHeader1 = _Structure(  # `version` field pulled out into VariablesRecord
137     name='VarHeader1',
138     fields=[
139         _Field('h', 'numSysVars', help='Number of system variables (K0, K1, ...).'),
140         _Field('h', 'numUserVars', help='Number of user numeric variables -- may be zero.'),
141         _Field('h', 'numUserStrs', help='Number of user string variables -- may be zero.'),
142         ])
143
144 # From Variables.h
145 VarHeader2 = _Structure(  # `version` field pulled out into VariablesRecord
146     name='VarHeader2',
147     fields=[
148         _Field('h', 'numSysVars', help='Number of system variables (K0, K1, ...).'),
149         _Field('h', 'numUserVars', help='Number of user numeric variables -- may be zero.'),
150         _Field('h', 'numUserStrs', help='Number of user string variables -- may be zero.'),
151         _Field('h', 'numDependentVars', help='Number of dependent numeric variables -- may be zero.'),
152         _Field('h', 'numDependentStrs', help='Number of dependent string variables -- may be zero.'),
153         ])
154
155 # From Variables.h
156 UserStrVarRec1 = _DynamicStructure(
157     name='UserStrVarRec1',
158     fields=[
159         ListedStaticStringField('c', 'name', help='Name of the string variable.', count=32),
160         _Field('h', 'strLen', help='The real size of the following array.'),
161         ListedDynamicStrDataField('c', 'data'),
162         ])
163
164 # From Variables.h
165 UserStrVarRec2 = _DynamicStructure(
166     name='UserStrVarRec2',
167     fields=[
168         ListedStaticStringField('c', 'name', help='Name of the string variable.', count=32),
169         _Field('l', 'strLen', help='The real size of the following array.'),
170         _Field('c', 'data'),
171         ])
172
173 # From Variables.h
174 VarNumRec = _Structure(
175     name='VarNumRec',
176     fields=[
177         _Field('h', 'numType', help='Type from binarywave.TYPE_TABLE'),
178         _Field('d', 'realPart', help='The real part of the number.'),
179         _Field('d', 'imagPart', help='The imag part if the number is complex.'),
180         _Field('l', 'reserved', help='Reserved - set to zero.'),
181         ])
182
183 # From Variables.h
184 UserNumVarRec = _DynamicStructure(
185     name='UserNumVarRec',
186     fields=[
187         ListedStaticStringField('c', 'name', help='Name of the string variable.', count=32),
188         _Field('h', 'type', help='0 = string, 1 = numeric.'),
189         DynamicVarNumField(VarNumRec, 'num', help='Type and value of the variable if it is numeric.  Not used for string.'),
190         ])
191
192 # From Variables.h
193 UserDependentVarRec = _DynamicStructure(
194     name='UserDependentVarRec',
195     fields=[
196         ListedStaticStringField('c', 'name', help='Name of the string variable.', count=32),
197         _Field('h', 'type', help='0 = string, 1 = numeric.'),
198         _Field(VarNumRec, 'num', help='Type and value of the variable if it is numeric.  Not used for string.'),
199         _Field('h', 'formulaLen', help='The length of the dependency formula.'),
200         DynamicFormulaField('c', 'formula', help='Start of the dependency formula. A C string including null terminator.'),
201         ])
202
203
204 class DynamicVarHeaderField (_DynamicField):
205     def pre_pack(self, parents, data):
206         raise NotImplementedError()
207
208     def post_unpack(self, parents, data):
209         var_structure = parents[-1]
210         var_data = self._get_structure_data(
211             parents, data, var_structure)
212         var_header_structure = self.format
213         data = var_data['var_header']
214         sys_vars_field = var_structure.get_field('sysVars')
215         sys_vars_field.count = data['numSysVars']
216         sys_vars_field.setup()
217         user_vars_field = var_structure.get_field('userVars')
218         user_vars_field.count = data['numUserVars']
219         user_vars_field.setup()
220         user_strs_field = var_structure.get_field('userStrs')
221         user_strs_field.count = data['numUserStrs']
222         user_strs_field.setup()
223         if 'numDependentVars' in data:
224             dependent_vars_field = var_structure.get_field('dependentVars')
225             dependent_vars_field.count = data['numDependentVars']
226             dependent_vars_field.setup()
227             dependent_strs_field = var_structure.get_field('dependentStrs')
228             dependent_strs_field.count = data['numDependentStrs']
229             dependent_strs_field.setup()
230         var_structure.setup()
231
232
233 Variables1 = _DynamicStructure(
234     name='Variables1',
235     fields=[
236         DynamicVarHeaderField(VarHeader1, 'var_header', help='Variables header'),
237         DynamicSysVarField('f', 'sysVars', help='System variables', count=0),
238         DynamicUserVarField(UserNumVarRec, 'userVars', help='User numeric variables', count=0),
239         DynamicUserStrField(UserStrVarRec1, 'userStrs', help='User string variables', count=0),
240         ])
241
242
243 Variables2 = _DynamicStructure(
244     name='Variables2',
245     fields=[
246         DynamicVarHeaderField(VarHeader2, 'var_header', help='Variables header'),
247         DynamicSysVarField('f', 'sysVars', help='System variables', count=0),
248         DynamicUserVarField(UserNumVarRec, 'userVars', help='User numeric variables', count=0),
249         DynamicUserStrField(UserStrVarRec2, 'userStrs', help='User string variables', count=0),
250         _Field(UserDependentVarRec, 'dependentVars', help='Dependent numeric variables.', count=0),
251         _Field(UserDependentVarRec, 'dependentStrs', help='Dependent string variables.', count=0),
252         ])
253
254
255 class DynamicVersionField (_DynamicField):
256     def pre_pack(self, parents, byte_order):
257         raise NotImplementedError()
258
259     def post_unpack(self, parents, data):
260         variables_structure = parents[-1]
261         variables_data = self._get_structure_data(
262             parents, data, variables_structure)
263         version = variables_data['version']
264         if variables_structure.byte_order in '@=':
265             need_to_reorder_bytes = _need_to_reorder_bytes(version)
266             variables_structure.byte_order = _byte_order(need_to_reorder_bytes)
267             _LOG.debug(
268                 'get byte order from version: {} (reorder? {})'.format(
269                     variables_structure.byte_order, need_to_reorder_bytes))
270         else:
271             need_to_reorder_bytes = False
272
273         old_format = variables_structure.fields[-1].format
274         if version == 1:
275             variables_structure.fields[-1].format = Variables1
276         elif version == 2:
277             variables_structure.fields[-1].format = Variables2
278         elif not need_to_reorder_bytes:
279             raise ValueError(
280                 'invalid variables record version: {}'.format(version))
281
282         if variables_structure.fields[-1].format != old_format:
283             _LOG.debug('change variables record from {} to {}'.format(
284                     old_format, variables_structure.fields[-1].format))
285             variables_structure.setup()
286         elif need_to_reorder_bytes:
287             variables_structure.setup()
288
289         # we might need to unpack again with the new byte order
290         return need_to_reorder_bytes
291
292
293 VariablesRecordStructure = _DynamicStructure(
294     name='VariablesRecord',
295     fields=[
296         DynamicVersionField('h', 'version', help='Version number for this header.'),
297         _Field(Variables1, 'variables', help='The rest of the variables data.'),
298         ])
299
300
301 class VariablesRecord (Record):
302     def __init__(self, *args, **kwargs):
303         super(VariablesRecord, self).__init__(*args, **kwargs)
304         # self.header['version']  # record version always 0?
305         VariablesRecordStructure.byte_order = '='
306         VariablesRecordStructure.setup()
307         stream = _io.BytesIO(bytes(self.data))
308         self.variables = VariablesRecordStructure.unpack_stream(stream)
309         self.namespace = {}
310         for key,value in self.variables['variables'].items():
311             if key not in ['var_header']:
312                 _LOG.debug('update namespace {} with {} for {}'.format(
313                         self.namespace, value, key))
314                 self.namespace.update(value)