example20.py 2.16 KB
Newer Older
1
2
3
#!/usr/bin/env python
from __future__ import print_function

4
import unittest
5
import numpy as np
6
from example import (
7
    create_rec_simple, create_rec_packed, create_rec_nested, print_format_descriptors,
8
9
    print_rec_simple, print_rec_packed, print_rec_nested, print_dtypes, get_format_unbound,
    create_rec_partial, create_rec_partial_nested
10
)
11
12
13
14
15


def check_eq(arr, data, dtype):
    np.testing.assert_equal(arr, np.array(data, dtype=dtype))

16
17
18
unittest.TestCase().assertRaisesRegex(
    RuntimeError, 'unsupported buffer format', get_format_unbound)

19
print_format_descriptors()
20
print_dtypes()
21

Ivan Smirnov's avatar
Ivan Smirnov committed
22
23
24
25
simple_dtype = np.dtype({'names': ['x', 'y', 'z'],
                         'formats': ['?', 'u4', 'f4'],
                         'offsets': [0, 4, 8]})
packed_dtype = np.dtype([('x', '?'), ('y', 'u4'), ('z', 'f4')])
26

27
28
elements = [(False, 0, 0.0), (True, 1, 1.5), (False, 2, 3.0)]

Ivan Smirnov's avatar
Ivan Smirnov committed
29
for func, dtype in [(create_rec_simple, simple_dtype), (create_rec_packed, packed_dtype)]:
Ivan Smirnov's avatar
Ivan Smirnov committed
30
31
32
33
34
    arr = func(0)
    assert arr.dtype == dtype
    check_eq(arr, [], simple_dtype)
    check_eq(arr, [], packed_dtype)

Ivan Smirnov's avatar
Ivan Smirnov committed
35
36
    arr = func(3)
    assert arr.dtype == dtype
37
38
    check_eq(arr, elements, simple_dtype)
    check_eq(arr, elements, packed_dtype)
39

40
41
42
43
    if dtype == simple_dtype:
        print_rec_simple(arr)
    else:
        print_rec_packed(arr)
Ivan Smirnov's avatar
Ivan Smirnov committed
44

45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

arr = create_rec_partial(3)
print(arr.dtype)
partial_dtype = arr.dtype
assert '' not in arr.dtype.fields
assert partial_dtype.itemsize > simple_dtype.itemsize
check_eq(arr, elements, simple_dtype)
check_eq(arr, elements, packed_dtype)

arr = create_rec_partial_nested(3)
print(arr.dtype)
assert '' not in arr.dtype.fields
assert '' not in arr.dtype.fields['a'][0].fields
assert arr.dtype.itemsize > partial_dtype.itemsize
np.testing.assert_equal(arr['a'], create_rec_partial(3))

Ivan Smirnov's avatar
Ivan Smirnov committed
61
62
63
64
65
66
67
68
69
70
71
nested_dtype = np.dtype([('a', simple_dtype), ('b', packed_dtype)])

arr = create_rec_nested(0)
assert arr.dtype == nested_dtype
check_eq(arr, [], nested_dtype)

arr = create_rec_nested(3)
assert arr.dtype == nested_dtype
check_eq(arr, [((False, 0, 0.0), (True, 1, 1.5)),
               ((True, 1, 1.5), (False, 2, 3.0)),
               ((False, 2, 3.0), (True, 3, 4.5))], nested_dtype)
72
print_rec_nested(arr)