1from __future__ import absolute_import, print_function, division
2
3
4from petl.test.failonerror import test_failonerror
5from petl.test.helpers import ieq
6from petl.transform.conversions import convert, convertall, convertnumbers, \
7    replace, update, format, interpolate
8
9from functools import partial
10
11
12def test_convert():
13
14    table1 = (('foo', 'bar', 'baz'),
15              ('A', 1, 2),
16              ('B', '2', '3.4'),
17              (u'B', u'3', u'7.8', True),
18              ('D', 'xyz', 9.0),
19              ('E', None))
20
21    # test the simplest style - single field, lambda function
22    table2 = convert(table1, 'foo', lambda s: s.lower())
23    expect2 = (('foo', 'bar', 'baz'),
24               ('a', 1, 2),
25               ('b', '2', '3.4'),
26               (u'b', u'3', u'7.8', True),
27               ('d', 'xyz', 9.0),
28               ('e', None))
29    ieq(expect2, table2)
30    ieq(expect2, table2)
31
32    # test single field with method call
33    table3 = convert(table1, 'foo', 'lower')
34    expect3 = expect2
35    ieq(expect3, table3)
36
37    # test single field with method call with arguments
38    table4 = convert(table1, 'foo', 'replace', 'B', 'BB')
39    expect4 = (('foo', 'bar', 'baz'),
40               ('A', 1, 2),
41               ('BB', '2', '3.4'),
42               (u'BB', u'3', u'7.8', True),
43               ('D', 'xyz', 9.0),
44               ('E', None))
45    ieq(expect4, table4)
46
47    # test multiple fields with the same conversion
48    table5 = convert(table1, ('bar', 'baz'), str)
49    expect5 = (('foo', 'bar', 'baz'),
50               ('A', '1', '2'),
51               ('B', '2', '3.4'),
52               (u'B', u'3', u'7.8', True),
53               ('D', 'xyz', '9.0'),
54               ('E', 'None'))
55    ieq(expect5, table5)
56
57    # test convert with dictionary
58    table6 = convert(table1, 'foo', {'A': 'Z', 'B': 'Y'})
59    expect6 = (('foo', 'bar', 'baz'),
60               ('Z', 1, 2),
61               ('Y', '2', '3.4'),
62               (u'Y', u'3', u'7.8', True),
63               ('D', 'xyz', 9.0),
64               ('E', None))
65    ieq(expect6, table6)
66
67
68def test_convert_empty():
69    table = (('foo', 'bar'),)
70    expect = (('foo', 'bar'),)
71    actual = convert(table, 'foo', int)
72    ieq(expect, actual)
73
74
75def test_convert_indexes():
76
77    table1 = (('foo', 'bar', 'baz'),
78              ('A', 1, 2),
79              ('B', '2', '3.4'),
80              (u'B', u'3', u'7.8', True),
81              ('D', 'xyz', 9.0),
82              ('E', None))
83
84    # test the simplest style - single field, lambda function
85    table2 = convert(table1, 0, lambda s: s.lower())
86    expect2 = (('foo', 'bar', 'baz'),
87               ('a', 1, 2),
88               ('b', '2', '3.4'),
89               (u'b', u'3', u'7.8', True),
90               ('d', 'xyz', 9.0),
91               ('e', None))
92    ieq(expect2, table2)
93    ieq(expect2, table2)
94
95    # test single field with method call
96    table3 = convert(table1, 0, 'lower')
97    expect3 = expect2
98    ieq(expect3, table3)
99
100    # test single field with method call with arguments
101    table4 = convert(table1, 0, 'replace', 'B', 'BB')
102    expect4 = (('foo', 'bar', 'baz'),
103               ('A', 1, 2),
104               ('BB', '2', '3.4'),
105               (u'BB', u'3', u'7.8', True),
106               ('D', 'xyz', 9.0),
107               ('E', None))
108    ieq(expect4, table4)
109
110    # test multiple fields with the same conversion
111    table5a = convert(table1, (1, 2), str)
112    table5b = convert(table1, (1, 'baz'), str)
113    table5c = convert(table1, ('bar', 2), str)
114    table5d = convert(table1, list(range(1, 3)), str)
115    expect5 = (('foo', 'bar', 'baz'),
116               ('A', '1', '2'),
117               ('B', '2', '3.4'),
118               (u'B', u'3', u'7.8', True),
119               ('D', 'xyz', '9.0'),
120               ('E', 'None'))
121    ieq(expect5, table5a)
122    ieq(expect5, table5b)
123    ieq(expect5, table5c)
124    ieq(expect5, table5d)
125
126    # test convert with dictionary
127    table6 = convert(table1, 0, {'A': 'Z', 'B': 'Y'})
128    expect6 = (('foo', 'bar', 'baz'),
129               ('Z', 1, 2),
130               ('Y', '2', '3.4'),
131               (u'Y', u'3', u'7.8', True),
132               ('D', 'xyz', 9.0),
133               ('E', None))
134    ieq(expect6, table6)
135
136
137def test_fieldconvert():
138
139    table1 = (('foo', 'bar', 'baz'),
140              ('A', 1, 2),
141              ('B', '2', '3.4'),
142              (u'B', u'3', u'7.8', True),
143              ('D', 'xyz', 9.0),
144              ('E', None))
145
146    # test the style where the converters functions are passed in as a
147    # dictionary
148    converters = {'foo': str, 'bar': int, 'baz': float}
149    table5 = convert(table1, converters, errorvalue='error')
150    expect5 = (('foo', 'bar', 'baz'),
151               ('A', 1, 2.0),
152               ('B', 2, 3.4),
153               ('B', 3, 7.8, True),  # N.B., long rows are preserved
154               ('D', 'error', 9.0),
155               ('E', 'error'))  # N.B., short rows are preserved
156    ieq(expect5, table5)
157
158    # test the style where the converters functions are added one at a time
159    table6 = convert(table1, errorvalue='err')
160    table6['foo'] = str
161    table6['bar'] = int
162    table6['baz'] = float
163    expect6 = (('foo', 'bar', 'baz'),
164               ('A', 1, 2.0),
165               ('B', 2, 3.4),
166               ('B', 3, 7.8, True),
167               ('D', 'err', 9.0),
168               ('E', 'err'))
169    ieq(expect6, table6)
170
171    # test some different converters
172    table7 = convert(table1)
173    table7['foo'] = 'replace', 'B', 'BB'
174    expect7 = (('foo', 'bar', 'baz'),
175               ('A', 1, 2),
176               ('BB', '2', '3.4'),
177               (u'BB', u'3', u'7.8', True),
178               ('D', 'xyz', 9.0),
179               ('E', None))
180    ieq(expect7, table7)
181
182    # test the style where the converters functions are passed in as a list
183    converters = [str, int, float]
184    table8 = convert(table1, converters, errorvalue='error')
185    expect8 = (('foo', 'bar', 'baz'),
186               ('A', 1, 2.0),
187               ('B', 2, 3.4),
188               ('B', 3, 7.8, True),  # N.B., long rows are preserved
189               ('D', 'error', 9.0),
190               ('E', 'error'))  # N.B., short rows are preserved
191    ieq(expect8, table8)
192
193    # test the style where the converters functions are passed in as a list
194    converters = [str, None, float]
195    table9 = convert(table1, converters, errorvalue='error')
196    expect9 = (('foo', 'bar', 'baz'),
197               ('A', 1, 2.0),
198               ('B', '2', 3.4),
199               ('B', u'3', 7.8, True),  # N.B., long rows are preserved
200               ('D', 'xyz', 9.0),
201               ('E', None))  # N.B., short rows are preserved
202    ieq(expect9, table9)
203
204
205def test_convertall():
206
207    table1 = (('foo', 'bar', 'baz'),
208              ('1', '3', '9'),
209              ('2', '1', '7'))
210    table2 = convertall(table1, int)
211    expect2 = (('foo', 'bar', 'baz'),
212               (1, 3, 9),
213               (2, 1, 7))
214    ieq(expect2, table2)
215    ieq(expect2, table2)
216
217
218def test_convertnumbers():
219
220    table1 = (('foo', 'bar', 'baz', 'quux'),
221              ('1', '3.0', '9+3j', 'aaa'),
222              ('2', '1.3', '7+2j', None))
223    table2 = convertnumbers(table1)
224    expect2 = (('foo', 'bar', 'baz', 'quux'),
225               (1, 3.0, 9+3j, 'aaa'),
226               (2, 1.3, 7+2j, None))
227    ieq(expect2, table2)
228    ieq(expect2, table2)
229
230
231def test_convert_translate():
232
233    table = (('foo', 'bar'),
234             ('M', 12),
235             ('F', 34),
236             ('-', 56))
237
238    trans = {'M': 'male', 'F': 'female'}
239    result = convert(table, 'foo', trans)
240    expectation = (('foo', 'bar'),
241                   ('male', 12),
242                   ('female', 34),
243                   ('-', 56))
244    ieq(expectation, result)
245
246
247def test_convert_with_row():
248
249    table = (('foo', 'bar'),
250             ('a', 1),
251             ('b', 2))
252
253    expect = (('foo', 'bar'),
254              ('a', 'A'),
255              ('b', 'B'))
256
257    actual = convert(table, 'bar',
258                     lambda v, row: row.foo.upper(),
259                     pass_row=True)
260    ieq(expect, actual)
261
262
263def test_convert_with_row_backwards_compat():
264
265    table = (('foo', 'bar'),
266             (' a ', 1),
267             (' b ', 2))
268
269    expect = (('foo', 'bar'),
270              ('a', 1),
271              ('b', 2))
272
273    actual = convert(table, 'foo', 'strip')
274    ieq(expect, actual)
275
276
277def test_convert_where():
278
279    tbl1 = (('foo', 'bar'),
280            ('a', 1),
281            ('b', 2))
282
283    expect = (('foo', 'bar'),
284              ('a', 1),
285              ('b', 4))
286
287    actual = convert(tbl1, 'bar', lambda v: v*2, where=lambda r: r.foo == 'b')
288    ieq(expect, actual)
289    ieq(expect, actual)
290    actual = convert(tbl1, 'bar', lambda v: v*2, where="{foo} == 'b'")
291    ieq(expect, actual)
292    ieq(expect, actual)
293
294
295def test_convert_failonerror():
296    input_  = (('foo',), ('A',), (1,))
297    cvt_    = {'foo': 'lower'}
298    expect_ = (('foo',), ('a',), (None,))
299
300    test_failonerror(
301            input_fn=partial(convert, input_, cvt_),
302            expected_output=expect_)
303
304
305def test_replace_where():
306
307    tbl1 = (('foo', 'bar'),
308            ('a', 1),
309            ('b', 2))
310
311    expect = (('foo', 'bar'),
312              ('a', 1),
313              ('b', 4))
314
315    actual = replace(tbl1, 'bar', 2, 4, where=lambda r: r.foo == 'b')
316    ieq(expect, actual)
317    ieq(expect, actual)
318    actual = replace(tbl1, 'bar', 2, 4, where="{foo} == 'b'")
319    ieq(expect, actual)
320    ieq(expect, actual)
321
322
323def test_update():
324
325    table1 = (('foo', 'bar', 'baz'),
326              ('A', 1, 2),
327              ('B', '2', '3.4'),
328              (u'B', u'3', u'7.8', True),
329              ('D', 'xyz', 9.0),
330              ('E', None))
331
332    table2 = update(table1, 'foo', 'X')
333    expect2 = (('foo', 'bar', 'baz'),
334               ('X', 1, 2),
335               ('X', '2', '3.4'),
336               ('X', u'3', u'7.8', True),
337               ('X', 'xyz', 9.0),
338               ('X', None))
339    ieq(expect2, table2)
340    ieq(expect2, table2)
341
342
343def test_replace_unhashable():
344
345    table1 = (('foo', 'bar'), ('a', ['b']), ('c', None))
346    expect = (('foo', 'bar'), ('a', ['b']), ('c', []))
347    actual = replace(table1, 'bar', None, [])
348    ieq(expect, actual)
349
350
351def test_format():
352
353    table = (('foo', 'bar'),
354             ('a', 1),
355             ('b', 2))
356
357    expect = (('foo', 'bar'),
358              ('a', '01'),
359              ('b', '02'))
360
361    actual = format(table, 'bar', '{0:02d}')
362    ieq(expect, actual)
363    ieq(expect, actual)
364
365
366def test_interpolate():
367
368    table = (('foo', 'bar'),
369             ('a', 1),
370             ('b', 2))
371
372    expect = (('foo', 'bar'),
373              ('a', '01'),
374              ('b', '02'))
375
376    actual = interpolate(table, 'bar', '%02d')
377    ieq(expect, actual)
378    ieq(expect, actual)
379
380