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