File: writer.go

package info (click to toggle)
golang-github-jonas-p-go-shp 0.1.1-2
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, bullseye, sid, trixie
  • size: 324 kB
  • sloc: makefile: 3
file content (345 lines) | stat: -rw-r--r-- 9,804 bytes parent folder | download
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
package shp

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// Writer is the type that is used to write a new shapefile.
type Writer struct {
	filename     string
	shp          writeSeekCloser
	shx          writeSeekCloser
	GeometryType ShapeType
	num          int32
	bbox         Box

	dbf             writeSeekCloser
	dbfFields       []Field
	dbfHeaderLength int16
	dbfRecordLength int16
}

type writeSeekCloser interface {
	io.Writer
	io.Seeker
	io.Closer
}

// Create returns a point to new Writer and the first error that was
// encountered. In case an error occurred the returned Writer point will be nil
// This also creates a corresponding SHX file. It is important to use Close()
// when done because that method writes all the headers for each file (SHP, SHX
// and DBF).
// If filename does not end on ".shp" already, it will be treated as the basename
// for the file and the ".shp" extension will be appended to that name.
func Create(filename string, t ShapeType) (*Writer, error) {
	if strings.HasSuffix(strings.ToLower(filename), ".shp") {
		filename = filename[0 : len(filename)-4]
	}
	shp, err := os.Create(filename + ".shp")
	if err != nil {
		return nil, err
	}
	shx, err := os.Create(filename + ".shx")
	if err != nil {
		return nil, err
	}
	shp.Seek(100, io.SeekStart)
	shx.Seek(100, io.SeekStart)
	w := &Writer{
		filename:     filename,
		shp:          shp,
		shx:          shx,
		GeometryType: t,
	}
	return w, nil
}

// Append returns a Writer pointer that will append to the given shapefile and
// the first error that was encounted during creation of that Writer. The
// shapefile must have a valid index file.
func Append(filename string) (*Writer, error) {
	shp, err := os.OpenFile(filename, os.O_RDWR, 0666)
	if err != nil {
		return nil, err
	}
	ext := filepath.Ext(filename)
	basename := filename[:len(filename)-len(ext)]
	w := &Writer{
		filename: basename,
		shp:      shp,
	}
	_, err = shp.Seek(32, io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHP geometry type: %v", err)
	}
	err = binary.Read(shp, binary.LittleEndian, &w.GeometryType)
	if err != nil {
		return nil, fmt.Errorf("cannot read geometry type: %v", err)
	}
	er := &errReader{Reader: shp}
	w.bbox.MinX = readFloat64(er)
	w.bbox.MinY = readFloat64(er)
	w.bbox.MaxX = readFloat64(er)
	w.bbox.MaxY = readFloat64(er)
	if er.e != nil {
		return nil, fmt.Errorf("cannot read bounding box: %v", er.e)
	}

	shx, err := os.OpenFile(basename+".shx", os.O_RDWR, 0666)
	if os.IsNotExist(err) {
		// TODO allow index file to not exist, in that case just
		// read through all the shapes and create it on the fly
	}
	if err != nil {
		return nil, fmt.Errorf("cannot open shapefile index: %v", err)
	}
	_, err = shx.Seek(-8, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to last shape index: %v", err)
	}
	var offset int32
	err = binary.Read(shx, binary.BigEndian, &offset)
	if err != nil {
		return nil, fmt.Errorf("cannot read last shape index: %v", err)
	}
	offset = offset * 2
	_, err = shp.Seek(int64(offset), io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to last shape: %v", err)
	}
	err = binary.Read(shp, binary.BigEndian, &w.num)
	if err != nil {
		return nil, fmt.Errorf("cannot read number of last shape: %v", err)
	}
	_, err = shp.Seek(0, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHP end: %v", err)
	}
	_, err = shx.Seek(0, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHX end: %v", err)
	}
	w.shx = shx

	dbf, err := os.Open(basename + ".dbf")
	if os.IsNotExist(err) {
		return w, nil // it's okay if the DBF does not exist
	}
	if err != nil {
		return nil, fmt.Errorf("cannot open DBF: %v", err)
	}

	_, err = dbf.Seek(8, io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek in DBF: %v", err)
	}
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfHeaderLength)
	if err != nil {
		return nil, fmt.Errorf("cannot read header length from DBF: %v", err)
	}
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfRecordLength)
	if err != nil {
		return nil, fmt.Errorf("cannot read record length from DBF: %v", err)
	}

	_, err = dbf.Seek(20, io.SeekCurrent) // skip padding
	if err != nil {
		return nil, fmt.Errorf("cannot seek in DBF: %v", err)
	}
	numFields := int(math.Floor(float64(w.dbfHeaderLength-33) / 32.0))
	w.dbfFields = make([]Field, numFields)
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfFields)
	if err != nil {
		return nil, fmt.Errorf("cannot read number of fields from DBF: %v", err)
	}
	_, err = dbf.Seek(0, io.SeekEnd) // skip padding
	if err != nil {
		return nil, fmt.Errorf("cannot seek to DBF end: %v", err)
	}
	w.dbf = dbf

	return w, nil
}

// Write shape to the Shapefile. This also creates
// a record in the SHX file and DBF file (if it is
// initialized). Returns the index of the written object
// which can be used in WriteAttribute.
func (w *Writer) Write(shape Shape) int32 {
	// increate bbox
	if w.num == 0 {
		w.bbox = shape.BBox()
	} else {
		w.bbox.Extend(shape.BBox())
	}

	w.num++
	binary.Write(w.shp, binary.BigEndian, w.num)
	w.shp.Seek(4, io.SeekCurrent)
	start, _ := w.shp.Seek(0, io.SeekCurrent)
	binary.Write(w.shp, binary.LittleEndian, w.GeometryType)
	shape.write(w.shp)
	finish, _ := w.shp.Seek(0, io.SeekCurrent)
	length := int32(math.Floor((float64(finish) - float64(start)) / 2.0))
	w.shp.Seek(start-4, io.SeekStart)
	binary.Write(w.shp, binary.BigEndian, length)
	w.shp.Seek(finish, io.SeekStart)

	// write shx
	binary.Write(w.shx, binary.BigEndian, int32((start-8)/2))
	binary.Write(w.shx, binary.BigEndian, length)

	// write empty record to dbf
	if w.dbf != nil {
		w.writeEmptyRecord()
	}

	return w.num - 1
}

// Close closes the Writer. This must be used at the end of
// the transaction because it writes the correct headers
// to the SHP/SHX and DBF files before closing.
func (w *Writer) Close() {
	w.writeHeader(w.shx)
	w.writeHeader(w.shp)
	w.shp.Close()
	w.shx.Close()

	if w.dbf == nil {
		w.SetFields([]Field{})
	}
	w.writeDbfHeader(w.dbf)
	w.dbf.Close()
}

// writeHeader wrires SHP/SHX headers to ws.
func (w *Writer) writeHeader(ws io.WriteSeeker) {
	filelength, _ := ws.Seek(0, io.SeekEnd)
	if filelength == 0 {
		filelength = 100
	}
	ws.Seek(0, io.SeekStart)
	// file code
	binary.Write(ws, binary.BigEndian, []int32{9994, 0, 0, 0, 0, 0})
	// file length
	binary.Write(ws, binary.BigEndian, int32(filelength/2))
	// version and shape type
	binary.Write(ws, binary.LittleEndian, []int32{1000, int32(w.GeometryType)})
	// bounding box
	binary.Write(ws, binary.LittleEndian, w.bbox)
	// elevation, measure
	binary.Write(ws, binary.LittleEndian, []float64{0.0, 0.0, 0.0, 0.0})
}

// writeDbfHeader writes a DBF header to ws.
func (w *Writer) writeDbfHeader(ws io.WriteSeeker) {
	ws.Seek(0, 0)
	// version, year (YEAR-1990), month, day
	binary.Write(ws, binary.LittleEndian, []byte{3, 24, 5, 3})
	// number of records
	binary.Write(ws, binary.LittleEndian, w.num)
	// header length, record length
	binary.Write(ws, binary.LittleEndian, []int16{w.dbfHeaderLength, w.dbfRecordLength})
	// padding
	binary.Write(ws, binary.LittleEndian, make([]byte, 20))

	for _, field := range w.dbfFields {
		binary.Write(ws, binary.LittleEndian, field)
	}

	// end with return
	ws.Write([]byte("\r"))
}

// SetFields sets field values in the DBF. This initializes the DBF file and
// should be used prior to writing any attributes.
func (w *Writer) SetFields(fields []Field) error {
	if w.dbf != nil {
		return errors.New("Cannot set fields in existing dbf")
	}

	var err error
	w.dbf, err = os.Create(w.filename + "dbf")
	if err != nil {
		return fmt.Errorf("Failed to open %s.dbf: %v", w.filename, err)
	}
	w.dbfFields = fields

	// calculate record length
	w.dbfRecordLength = int16(1)
	for _, field := range w.dbfFields {
		w.dbfRecordLength += int16(field.Size)
	}

	// header lengh
	w.dbfHeaderLength = int16(len(w.dbfFields)*32 + 33)

	// fill header space with empty bytes for now
	buf := make([]byte, w.dbfHeaderLength)
	binary.Write(w.dbf, binary.LittleEndian, buf)

	// write empty records
	for n := int32(0); n < w.num; n++ {
		w.writeEmptyRecord()
	}
	return nil
}

// Writes an empty record to the end of the DBF. This
// works by seeking to the end of the file and writing
// dbfRecordLength number of bytes. The first byte is a
// space that indicates a new record.
func (w *Writer) writeEmptyRecord() {
	w.dbf.Seek(0, io.SeekEnd)
	buf := make([]byte, w.dbfRecordLength)
	buf[0] = ' '
	binary.Write(w.dbf, binary.LittleEndian, buf)
}

// WriteAttribute writes value for field into the given row in the DBF. Row
// number should be the same as the order the Shape was written to the
// Shapefile. The field value corresponds to the field in the slice used in
// SetFields.
func (w *Writer) WriteAttribute(row int, field int, value interface{}) error {
	var buf []byte
	switch v := value.(type) {
	case int:
		buf = []byte(strconv.Itoa(v))
	case float64:
		precision := w.dbfFields[field].Precision
		buf = []byte(strconv.FormatFloat(v, 'f', int(precision), 64))
	case string:
		buf = []byte(v)
	default:
		return fmt.Errorf("Unsupported value type: %T", v)
	}

	if w.dbf == nil {
		return errors.New("Initialize DBF by using SetFields first")
	}
	if sz := int(w.dbfFields[field].Size); len(buf) > sz {
		return fmt.Errorf("Unable to write field %v: %q exceeds field length %v", field, buf, sz)
	}

	seekTo := 1 + int64(w.dbfHeaderLength) + (int64(row) * int64(w.dbfRecordLength))
	for n := 0; n < field; n++ {
		seekTo += int64(w.dbfFields[n].Size)
	}
	w.dbf.Seek(seekTo, io.SeekStart)
	return binary.Write(w.dbf, binary.LittleEndian, buf)
}

// BBox returns the bounding box of the Writer.
func (w *Writer) BBox() Box {
	return w.bbox
}