-
Notifications
You must be signed in to change notification settings - Fork 214
/
test_clib_put_strings.py
64 lines (59 loc) · 2.15 KB
/
test_clib_put_strings.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
"""
Test the functions that put string data into GMT.
"""
import numpy as np
import numpy.testing as npt
import pytest
from pygmt import clib
from pygmt.exceptions import GMTCLibError
from pygmt.helpers import GMTTempFile
@pytest.mark.benchmark
def test_put_strings():
"""
Check that assigning a numpy array of dtype str to a dataset works.
"""
with clib.Session() as lib:
dataset = lib.create_data(
family="GMT_IS_DATASET|GMT_VIA_VECTOR",
geometry="GMT_IS_POINT",
mode="GMT_CONTAINER_ONLY",
dim=[2, 5, 1, 0], # columns, rows, layers, dtype
)
x = np.array([1, 2, 3, 4, 5], dtype=np.int32)
y = np.array([6, 7, 8, 9, 10], dtype=np.int32)
strings = np.array(["a", "bc", "defg", "hijklmn", "opqrst"], dtype=str)
lib.put_vector(dataset, column=lib["GMT_X"], vector=x)
lib.put_vector(dataset, column=lib["GMT_Y"], vector=y)
lib.put_strings(
dataset, family="GMT_IS_VECTOR|GMT_IS_DUPLICATE", strings=strings
)
# Turns out wesn doesn't matter for Datasets
wesn = [0] * 6
# Save the data to a file to see if it's being accessed correctly
with GMTTempFile() as tmp_file:
lib.write_data(
"GMT_IS_VECTOR",
"GMT_IS_POINT",
"GMT_WRITE_SET",
wesn,
tmp_file.name,
dataset,
)
# Load the data and check that it's correct
newx, newy, newstrings = tmp_file.loadtxt(
unpack=True, dtype=[("x", np.int32), ("y", np.int32), ("text", "<U7")]
)
npt.assert_array_equal(newx, x)
npt.assert_array_equal(newy, y)
npt.assert_array_equal(newstrings, strings)
def test_put_strings_fails():
"""
Check that put_strings raises an exception if return code is not zero.
"""
with clib.Session() as lib:
with pytest.raises(GMTCLibError):
lib.put_strings(
dataset=None,
family="GMT_IS_VECTOR|GMT_IS_DUPLICATE",
strings=np.empty(shape=(3,), dtype=str),
)