writer.go

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
}