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 }