1import unittest
2import sys
3import io
4
5# import pysal_examples
6from .... import examples as pysal_examples
7
8from ..shapefile import (
9    noneMax,
10    noneMin,
11    shp_file,
12    shx_file,
13    NullShape,
14    Point,
15    PolyLine,
16    MultiPoint,
17    PointZ,
18    PolyLineZ,
19    PolygonZ,
20    MultiPointZ,
21    PointM,
22    PolyLineM,
23    PolygonM,
24    MultiPointM,
25    MultiPatch,
26)
27import os
28
29
30def bufferIO(buf):
31    """Temp stringIO function to force compat.
32    """
33    return io.BytesIO(buf)
34
35
36class TestNoneMax(unittest.TestCase):
37    def test_none_max(self):
38        self.assertEqual(5, noneMax(5, None))
39        self.assertEqual(1, noneMax(None, 1))
40        self.assertEqual(None, noneMax(None, None))
41
42
43class TestNoneMin(unittest.TestCase):
44    def test_none_min(self):
45        self.assertEqual(5, noneMin(5, None))
46        self.assertEqual(1, noneMin(None, 1))
47        self.assertEqual(None, noneMin(None, None))
48
49
50class test_shp_file(unittest.TestCase):
51    def test___init__(self):
52        shp = shp_file(pysal_examples.get_path("10740.shp"))
53        assert shp.header == {
54            "BBOX Xmax": -105.29012,
55            "BBOX Ymax": 36.219799000000002,
56            "BBOX Mmax": 0.0,
57            "BBOX Zmin": 0.0,
58            "BBOX Mmin": 0.0,
59            "File Code": 9994,
60            "BBOX Ymin": 34.259672000000002,
61            "BBOX Xmin": -107.62651,
62            "Unused0": 0,
63            "Unused1": 0,
64            "Unused2": 0,
65            "Unused3": 0,
66            "Unused4": 0,
67            "Version": 1000,
68            "BBOX Zmax": 0.0,
69            "Shape Type": 5,
70            "File Length": 260534,
71        }
72
73    def test___iter__(self):
74        shp = shp_file(pysal_examples.get_path("Point.shp"))
75        points = [pt for pt in shp]
76        expected = [
77            {"Y": -0.25904661905760773, "X": -0.00068176617532103578, "Shape Type": 1},
78            {"Y": -0.25630328607387354, "X": 0.11697145363360706, "Shape Type": 1},
79            {"Y": -0.33930131004366804, "X": 0.05043668122270728, "Shape Type": 1},
80            {"Y": -0.41266375545851519, "X": -0.041266375545851552, "Shape Type": 1},
81            {"Y": -0.44017467248908293, "X": -0.011462882096069604, "Shape Type": 1},
82            {"Y": -0.46080786026200882, "X": 0.027510917030567628, "Shape Type": 1},
83            {"Y": -0.45851528384279472, "X": 0.075655021834060809, "Shape Type": 1},
84            {"Y": -0.43558951965065495, "X": 0.11233624454148461, "Shape Type": 1},
85            {"Y": -0.40578602620087334, "X": 0.13984716157205224, "Shape Type": 1},
86        ]
87        assert points == expected
88
89    def test___len__(self):
90        shp = shp_file(pysal_examples.get_path("10740.shp"))
91        assert len(shp) == 195
92
93    def test_add_shape(self):
94        shp = shp_file("test_point", "w", "POINT")
95        points = [
96            {"Shape Type": 1, "X": 0, "Y": 0},
97            {"Shape Type": 1, "X": 1, "Y": 1},
98            {"Shape Type": 1, "X": 2, "Y": 2},
99            {"Shape Type": 1, "X": 3, "Y": 3},
100            {"Shape Type": 1, "X": 4, "Y": 4},
101        ]
102        for pt in points:
103            shp.add_shape(pt)
104        shp.close()
105
106        for a, b in zip(points, shp_file("test_point")):
107            self.assertEqual(a, b)
108        os.remove("test_point.shp")
109        os.remove("test_point.shx")
110
111    def test_close(self):
112        shp = shp_file(pysal_examples.get_path("10740.shp"))
113        shp.close()
114        self.assertEqual(shp.fileObj.closed, True)
115
116    def test_get_shape(self):
117        shp = shp_file(pysal_examples.get_path("Line.shp"))
118        rec = shp.get_shape(0)
119        expected = {
120            "BBOX Ymax": -0.25832280562918325,
121            "NumPoints": 3,
122            "BBOX Ymin": -0.25895877033237352,
123            "NumParts": 1,
124            "Vertices": [
125                (-0.0090539248870159517, -0.25832280562918325),
126                (0.0074811573959305822, -0.25895877033237352),
127                (0.0074811573959305822, -0.25895877033237352),
128            ],
129            "BBOX Xmax": 0.0074811573959305822,
130            "BBOX Xmin": -0.0090539248870159517,
131            "Shape Type": 3,
132            "Parts Index": [0],
133        }
134        self.assertEqual(expected, shp.get_shape(0))
135
136    def test_next(self):
137        shp = shp_file(pysal_examples.get_path("Point.shp"))
138        points = [pt for pt in shp]
139        expected = {
140            "Y": -0.25904661905760773,
141            "X": -0.00068176617532103578,
142            "Shape Type": 1,
143        }
144        self.assertEqual(expected, next(shp))
145        expected = {
146            "Y": -0.25630328607387354,
147            "X": 0.11697145363360706,
148            "Shape Type": 1,
149        }
150        self.assertEqual(expected, next(shp))
151
152    def test_type(self):
153        shp = shp_file(pysal_examples.get_path("Point.shp"))
154        self.assertEqual("POINT", shp.type())
155        shp = shp_file(pysal_examples.get_path("Polygon.shp"))
156        self.assertEqual("POLYGON", shp.type())
157        shp = shp_file(pysal_examples.get_path("Line.shp"))
158        self.assertEqual("ARC", shp.type())
159
160
161class test_shx_file(unittest.TestCase):
162    def test___init__(self):
163        shx = shx_file(pysal_examples.get_path("Point.shx"))
164        assert isinstance(shx, shx_file)
165
166    def test_add_record(self):
167        shx = shx_file(pysal_examples.get_path("Point.shx"))
168        expectedIndex = [
169            (100, 20),
170            (128, 20),
171            (156, 20),
172            (184, 20),
173            (212, 20),
174            (240, 20),
175            (268, 20),
176            (296, 20),
177            (324, 20),
178        ]
179        assert shx.index == expectedIndex
180        shx2 = shx_file("test", "w")
181        for i, rec in enumerate(shx.index):
182            id, location = shx2.add_record(rec[1])
183            assert id == (i + 1)
184            assert location == rec[0]
185        assert shx2.index == shx.index
186        shx2.close(shx._header)
187        new_shx = open("test.shx", "rb").read()
188        expected_shx = open(pysal_examples.get_path("Point.shx"), "rb").read()
189        assert new_shx == expected_shx
190        os.remove("test.shx")
191
192    def test_close(self):
193        shx = shx_file(pysal_examples.get_path("Point.shx"))
194        shx.close(None)
195        self.assertEqual(shx.fileObj.closed, True)
196
197
198class TestNullShape(unittest.TestCase):
199    def test_pack(self):
200        null_shape = NullShape()
201        self.assertEqual(b"\x00" * 4, null_shape.pack())
202
203    def test_unpack(self):
204        null_shape = NullShape()
205        self.assertEqual(None, null_shape.unpack())
206
207
208class TestPoint(unittest.TestCase):
209    def test_pack(self):
210        record = {"X": 5, "Y": 5, "Shape Type": 1}
211        expected = b"\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x14\x40\x00\x00\x00\x00\x00\x00\x14\x40"
212        self.assertEqual(expected, Point.pack(record))
213
214    def test_unpack(self):
215        dat = bufferIO(
216            b"\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x14\x40\x00\x00\x00\x00\x00\x00\x14\x40"
217        )
218        expected = {"X": 5, "Y": 5, "Shape Type": 1}
219        self.assertEqual(expected, Point.unpack(dat))
220
221
222class TestPolyLine(unittest.TestCase):
223    def test_pack(self):
224        record = {
225            "BBOX Ymax": -0.25832280562918325,
226            "NumPoints": 3,
227            "BBOX Ymin": -0.25895877033237352,
228            "NumParts": 1,
229            "Vertices": [
230                (-0.0090539248870159517, -0.25832280562918325),
231                (0.0074811573959305822, -0.25895877033237352),
232                (0.0074811573959305822, -0.25895877033237352),
233            ],
234            "BBOX Xmax": 0.0074811573959305822,
235            "BBOX Xmin": -0.0090539248870159517,
236            "Shape Type": 3,
237            "Parts Index": [0],
238        }
239        expected = b"""\x03\x00\x00\x00\xc0\x46\x52\x3a\xdd\x8a\x82\
240\xbf\x3d\xc1\x65\xce\xc7\x92\xd0\xbf\x00\xc5\
241\xa0\xe5\x8f\xa4\x7e\x3f\x6b\x40\x7f\x60\x5c\
242\x88\xd0\xbf\x01\x00\x00\x00\x03\x00\x00\x00\
243\x00\x00\x00\x00\xc0\x46\x52\x3a\xdd\x8a\x82\
244\xbf\x6b\x40\x7f\x60\x5c\x88\xd0\xbf\x00\xc5\
245\xa0\xe5\x8f\xa4\x7e\x3f\x3d\xc1\x65\xce\xc7\
246\x92\xd0\xbf\x00\xc5\xa0\xe5\x8f\xa4\x7e\x3f\
247\x3d\xc1\x65\xce\xc7\x92\xd0\xbf"""
248        self.assertEqual(expected, PolyLine.pack(record))
249
250    def test_unpack(self):
251        dat = bufferIO(
252            b"""\x03\x00\x00\x00\xc0\x46\x52\x3a\xdd\x8a\x82\
253\xbf\x3d\xc1\x65\xce\xc7\x92\xd0\xbf\x00\xc5\
254\xa0\xe5\x8f\xa4\x7e\x3f\x6b\x40\x7f\x60\x5c\
255\x88\xd0\xbf\x01\x00\x00\x00\x03\x00\x00\x00\
256\x00\x00\x00\x00\xc0\x46\x52\x3a\xdd\x8a\x82\
257\xbf\x6b\x40\x7f\x60\x5c\x88\xd0\xbf\x00\xc5\
258\xa0\xe5\x8f\xa4\x7e\x3f\x3d\xc1\x65\xce\xc7\
259\x92\xd0\xbf\x00\xc5\xa0\xe5\x8f\xa4\x7e\x3f\
260\x3d\xc1\x65\xce\xc7\x92\xd0\xbf"""
261        )
262        expected = {
263            "BBOX Ymax": -0.25832280562918325,
264            "NumPoints": 3,
265            "BBOX Ymin": -0.25895877033237352,
266            "NumParts": 1,
267            "Vertices": [
268                (-0.0090539248870159517, -0.25832280562918325),
269                (0.0074811573959305822, -0.25895877033237352),
270                (0.0074811573959305822, -0.25895877033237352),
271            ],
272            "BBOX Xmax": 0.0074811573959305822,
273            "BBOX Xmin": -0.0090539248870159517,
274            "Shape Type": 3,
275            "Parts Index": [0],
276        }
277        self.assertEqual(expected, PolyLine.unpack(dat))
278
279
280class TestMultiPoint(unittest.TestCase):
281    def test___init__(self):
282        self.assertRaises(NotImplementedError, MultiPoint)
283
284
285class TestPointZ(unittest.TestCase):
286    def test_pack(self):
287        record = {"X": 5, "Y": 5, "Z": 5, "M": 5, "Shape Type": 11}
288        expected = b"\x0b\x00\x00\x00\x00\x00\x00\x00\x00\x00\x14@\x00\x00\x00\x00\x00\x00\x14@\x00\x00\x00\x00\x00\x00\x14@\x00\x00\x00\x00\x00\x00\x14@"
289        self.assertEqual(expected, PointZ.pack(record))
290
291    def test_unpack(self):
292        dat = bufferIO(
293            b"\x0b\x00\x00\x00\x00\x00\x00\x00\x00\x00\x14@\x00\x00\x00\x00\x00\x00\x14@\x00\x00\x00\x00\x00\x00\x14@\x00\x00\x00\x00\x00\x00\x14@"
294        )
295        expected = {"X": 5, "Y": 5, "Z": 5, "M": 5, "Shape Type": 11}
296        self.assertEqual(expected, PointZ.unpack(dat))
297
298
299class TestPolyLineZ(unittest.TestCase):
300    def test___init__(self):
301        self.assertRaises(NotImplementedError, PolyLineZ)
302
303
304class TestPolyLineZ(unittest.TestCase):
305    def test_pack(self):
306        record = {
307            "BBOX Ymax": -0.25832280562918325,
308            "NumPoints": 3,
309            "BBOX Ymin": -0.25895877033237352,
310            "NumParts": 1,
311            "Vertices": [
312                (-0.0090539248870159517, -0.25832280562918325),
313                (0.0074811573959305822, -0.25895877033237352),
314                (0.0074811573959305822, -0.25895877033237352),
315            ],
316            "BBOX Xmax": 0.0074811573959305822,
317            "BBOX Xmin": -0.0090539248870159517,
318            "Shape Type": 13,
319            "Parts Index": [0],
320            "Zmin": 0,
321            "Zmax": 10,
322            "Zarray": [0, 5, 10],
323            "Mmin": 2,
324            "Mmax": 4,
325            "Marray": [2, 3, 4],
326        }
327        expected = b"""\r\x00\x00\x00\xc0FR:\xdd\x8a\x82\xbf=\xc1e\xce\xc7\x92\xd0\xbf\x00\xc5\xa0\xe5\x8f\xa4~?k@\x7f`\\\x88\xd0\xbf\x01\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\xc0FR:\xdd\x8a\x82\xbfk@\x7f`\\\x88\xd0\xbf\x00\xc5\xa0\xe5\x8f\xa4~?=\xc1e\xce\xc7\x92\xd0\xbf\x00\xc5\xa0\xe5\x8f\xa4~?=\xc1e\xce\xc7\x92\xd0\xbf\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00$@\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x14@\x00\x00\x00\x00\x00\x00$@\x00\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x10@\x00\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x08@\x00\x00\x00\x00\x00\x00\x10@"""
328        self.assertEqual(expected, PolyLineZ.pack(record))
329
330    def test_unpack(self):
331        dat = bufferIO(
332            b"""\r\x00\x00\x00\xc0FR:\xdd\x8a\x82\xbf=\xc1e\xce\xc7\x92\xd0\xbf\x00\xc5\xa0\xe5\x8f\xa4~?k@\x7f`\\\x88\xd0\xbf\x01\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\xc0FR:\xdd\x8a\x82\xbfk@\x7f`\\\x88\xd0\xbf\x00\xc5\xa0\xe5\x8f\xa4~?=\xc1e\xce\xc7\x92\xd0\xbf\x00\xc5\xa0\xe5\x8f\xa4~?=\xc1e\xce\xc7\x92\xd0\xbf\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00$@\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x14@\x00\x00\x00\x00\x00\x00$@\x00\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x10@\x00\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x08@\x00\x00\x00\x00\x00\x00\x10@"""
333        )
334        expected = {
335            "BBOX Ymax": -0.25832280562918325,
336            "NumPoints": 3,
337            "BBOX Ymin": -0.25895877033237352,
338            "NumParts": 1,
339            "Vertices": [
340                (-0.0090539248870159517, -0.25832280562918325),
341                (0.0074811573959305822, -0.25895877033237352),
342                (0.0074811573959305822, -0.25895877033237352),
343            ],
344            "BBOX Xmax": 0.0074811573959305822,
345            "BBOX Xmin": -0.0090539248870159517,
346            "Shape Type": 13,
347            "Parts Index": [0],
348            "Zmin": 0,
349            "Zmax": 10,
350            "Zarray": [0, 5, 10],
351            "Mmin": 2,
352            "Mmax": 4,
353            "Marray": [2, 3, 4],
354        }
355        self.assertEqual(expected, PolyLineZ.unpack(dat))
356
357
358class TestPolygonZ(unittest.TestCase):
359    def test_pack(self):
360        record = {
361            "BBOX Xmin": 0.0,
362            "BBOX Xmax": 10.0,
363            "BBOX Ymin": 0.0,
364            "BBOX Ymax": 10.0,
365            "NumPoints": 4,
366            "NumParts": 1,
367            "Vertices": [(0.0, 0.0), (10.0, 10.0), (10.0, 0.0), (0.0, 0.0)],
368            "Shape Type": 15,
369            "Parts Index": [0],
370            "Zmin": 0,
371            "Zmax": 10,
372            "Zarray": [0, 10, 0, 0],
373            "Mmin": 2,
374            "Mmax": 4,
375            "Marray": [2, 4, 2, 2],
376        }
377        dat = bufferIO(PolygonZ.pack(record))
378        self.assertEqual(record, PolygonZ.unpack(dat))
379
380
381class TestMultiPointZ(unittest.TestCase):
382    def test___init__(self):
383        self.assertRaises(NotImplementedError, MultiPointZ)
384        # multi_point_z = MultiPointZ()
385
386
387class TestPointM(unittest.TestCase):
388    def test___init__(self):
389        self.assertRaises(NotImplementedError, PointM)
390        # point_m = PointM()
391
392
393class TestPolyLineM(unittest.TestCase):
394    def test___init__(self):
395        self.assertRaises(NotImplementedError, PolyLineM)
396        # poly_line_m = PolyLineM()
397
398
399class TestPolygonM(unittest.TestCase):
400    def test___init__(self):
401        self.assertRaises(NotImplementedError, PolygonM)
402        # polygon_m = PolygonM()
403
404
405class TestMultiPointM(unittest.TestCase):
406    def test___init__(self):
407        self.assertRaises(NotImplementedError, MultiPointM)
408        # multi_point_m = MultiPointM()
409
410
411class TestMultiPatch(unittest.TestCase):
412    def test___init__(self):
413        self.assertRaises(NotImplementedError, MultiPatch)
414        # multi_patch = MultiPatch()
415
416
417class _TestPoints(unittest.TestCase):
418    def test1(self):
419        """Test creating and reading Point Shape Files."""
420        shp = shp_file("test_point", "w", "POINT")
421        points = [
422            {"Shape Type": 1, "X": 0, "Y": 0},
423            {"Shape Type": 1, "X": 1, "Y": 1},
424            {"Shape Type": 1, "X": 2, "Y": 2},
425            {"Shape Type": 1, "X": 3, "Y": 3},
426            {"Shape Type": 1, "X": 4, "Y": 4},
427        ]
428        for pt in points:
429            shp.add_shape(pt)
430        shp.close()
431
432        shp = list(shp_file("test_point"))
433        for a, b in zip(points, shp):
434            self.assertEqual(a, b)
435        os.remove("test_point.shp")
436        os.remove("test_point.shx")
437
438
439class _TestPolyLines(unittest.TestCase):
440    def test1(self):
441        """Test creating and reading PolyLine Shape Files."""
442        lines = [[(0, 0), (4, 4)], [(1, 0), (5, 4)], [(2, 0), (6, 4)]]
443        shapes = []
444        for line in lines:
445            x = [v[0] for v in line]
446            y = [v[1] for v in line]
447            rec = {}
448            rec["BBOX Xmin"] = min(x)
449            rec["BBOX Ymin"] = min(y)
450            rec["BBOX Xmax"] = max(x)
451            rec["BBOX Ymax"] = max(y)
452            rec["NumPoints"] = len(line)
453            rec["NumParts"] = 1
454            rec["Vertices"] = line
455            rec["Shape Type"] = 3
456            rec["Parts Index"] = [0]
457            shapes.append(rec)
458        shp = shp_file("test_line", "w", "ARC")
459        for line in shapes:
460            shp.add_shape(line)
461        shp.close()
462        shp = list(shp_file("test_line"))
463        for a, b in zip(shapes, shp):
464            self.assertEqual(a, b)
465        os.remove("test_line.shp")
466        os.remove("test_line.shx")
467
468
469class _TestPolygons(unittest.TestCase):
470    def test1(self):
471        """Test creating and reading PolyLine Shape Files."""
472        lines = [
473            [(0, 0), (4, 4), (5, 4), (1, 0), (0, 0)],
474            [(1, 0), (5, 4), (6, 4), (2, 0), (1, 0)],
475        ]
476        shapes = []
477        for line in lines:
478            x = [v[0] for v in line]
479            y = [v[1] for v in line]
480            rec = {}
481            rec["BBOX Xmin"] = min(x)
482            rec["BBOX Ymin"] = min(y)
483            rec["BBOX Xmax"] = max(x)
484            rec["BBOX Ymax"] = max(y)
485            rec["NumPoints"] = len(line)
486            rec["NumParts"] = 1
487            rec["Vertices"] = line
488            rec["Shape Type"] = 5
489            rec["Parts Index"] = [0]
490            shapes.append(rec)
491        shp = shp_file("test_poly", "w", "POLYGON")
492        for line in shapes:
493            shp.add_shape(line)
494        shp.close()
495        shp = list(shp_file("test_poly"))
496        for a, b in zip(shapes, shp):
497            self.assertEqual(a, b)
498        os.remove("test_poly.shp")
499        os.remove("test_poly.shx")
500
501
502if __name__ == "__main__":
503    unittest.main()
504