Write extra columns in csv output

copier/copier.go

@@ -344,27 +344,39 @@ func (copier *Copier) checkEntity(ent tl.Entity) error {
 		}
 	}
 
+	// UpdateKeys is handled separately from other validators.
+	var referr error
+	if extEnt, ok := ent.(tl.EntityWithReferences); ok {
+		referr = extEnt.UpdateKeys(copier.EntityMap)
+	}
+
 	// Run Entity Validators
+	var errs []error
+	var warns []error
 	for _, v := range copier.errorValidators {
 		for _, err := range v.Validate(ent) {
-			ent.AddError(err)
+			errs = append(errs, err)
 		}
 	}
 	for _, v := range copier.warningValidators {
 		for _, err := range v.Validate(ent) {
-			ent.AddWarning(err)
+			warns = append(warns, err)
 		}
 	}
-
-	// UpdateKeys is handled separately from other validators.
-	referr := ent.UpdateKeys(copier.EntityMap)
-	if referr != nil {
-		ent.AddError(referr)
+	if extEnt, ok := ent.(tl.EntityWithErrors); ok {
+		for _, err := range errs {
+			extEnt.AddError(err)
+		}
+		for _, err := range warns {
+			extEnt.AddWarning(err)
+		}
+		if referr != nil {
+			extEnt.AddError(referr)
+		}
+		// Update to include the errors from entity validators
+		errs = extEnt.Errors()
+		warns = extEnt.Warnings()
 	}
-
-	// Perform entity level validation; includes any previous errors
-	errs := ent.Errors()
-	warns := ent.Warnings()
 	for _, err := range warns {
 		copier.sublogger.Debug().Str("filename", efn).Str("source_id", sid).Str("cause", err.Error()).Msg("warning")
 	}
@@ -378,7 +390,6 @@ func (copier *Copier) checkEntity(ent tl.Entity) error {
 		copier.result.SkipEntityErrorCount[efn]++
 		return errs[0]
 	}
-
 	if referr != nil && !copier.AllowReferenceErrors {
 		copier.result.SkipEntityReferenceCount[efn]++
 		return referr

internal/testutil/copy.go

@@ -17,8 +17,10 @@ func DirectCopy(reader tl.Reader, writer tl.Writer) error {
 	cp := func(ent tl.Entity) {
 		// All other entities
 		sid := ent.EntityID()
-		if err := ent.UpdateKeys(emap); err != nil {
-			errs = append(errs, fmt.Errorf("entity: %#v error: %s", ent, err))
+		if extEnt, ok := ent.(tl.EntityWithReferences); ok {
+			if err := extEnt.UpdateKeys(emap); err != nil {
+				errs = append(errs, fmt.Errorf("entity: %#v error: %s", ent, err))
+			}
 		}
 		eid, err := writer.AddEntity(ent)
 		if err != nil {

internal/testutil/expect.go

@@ -17,9 +17,12 @@ type HasContext interface {
 
 // GetExpectErrors gets any ExpectError specified by an Entity.
 func GetExpectErrors(ent tl.Entity) []ExpectError {
+	extEnt, ok := ent.(tl.EntityWithExtra)
+	if !ok {
+		return nil
+	}
 	ret := []ExpectError{}
-	ex := ent.Extra()
-	value, ok := ex["expect_error"]
+	value, ok := extEnt.GetExtra("expect_error")
 	if len(value) == 0 || !ok {
 		return ret
 	}

tl/gtfs_entity.go

@@ -5,71 +5,106 @@ import (
 	"time"
 )
 
-// EntityError is an interface for GTFS Errors
-type EntityError interface {
-	Error() string
-}
-
 // Entity provides an interface for GTFS entities.
 type Entity interface {
 	EntityID() string
 	Filename() string
+}
+
+type EntityWithReferences interface {
+	UpdateKeys(*EntityMap) error
+}
+
+type EntityWithExtra interface {
+	SetExtra(string, string)
+	GetExtra(string) (string, bool)
+	ExtraKeys() []string
+}
+
+type EntityWithErrors interface {
 	Errors() []error
 	Warnings() []error
 	AddError(error)
 	AddWarning(error)
-	SetExtra(string, string)
-	Extra() map[string]string
-	UpdateKeys(*EntityMap) error
 }
 
 /////////
 
-// MinEntity provides default methods.
+// MinEntity provides minimum set of default methods.
 type MinEntity struct {
-	extra        []string
-	loadErrors   []error
-	loadWarnings []error
 }
 
-// Extra provides any additional fields that were present.
-func (ent *MinEntity) Extra() map[string]string {
-	ret := map[string]string{}
+// Filename returns the filename for this entity.
+func (ent MinEntity) Filename() string { return "" }
+
+// EntityID returns the entity ID.
+func (ent MinEntity) EntityID() string { return "" }
+
+/////////
+
+type ExtraEntity struct {
+	extra []string
+}
+
+// SetExtra adds a string key, value pair to the entity's extra fields.
+func (ent *ExtraEntity) SetExtra(key string, value string) {
+	if key == "" {
+		return
+	}
+	for i := 0; i < len(ent.extra); i += 2 {
+		if ent.extra[i] == key {
+			return
+		}
+	}
+	ent.extra = append(ent.extra, key, value)
+}
+
+func (ent *ExtraEntity) GetExtra(key string) (string, bool) {
+	for i := 0; i < len(ent.extra); i += 2 {
+		if ent.extra[i] == key {
+			return ent.extra[i+1], true
+		}
+	}
+	return "", false
+}
+
+func (ent *ExtraEntity) ExtraKeys() []string {
+	var ret []string
 	for i := 0; i < len(ent.extra); i += 2 {
-		ret[ent.extra[i]] = ent.extra[i+1]
+		ret = append(ret, ent.extra[i])
 	}
 	return ret
 }
 
-// SetExtra adds a string key, value pair to the entity's extra fields.
-func (ent *MinEntity) SetExtra(key string, value string) {
-	ent.extra = append(ent.extra, key, value)
+/////////
+
+type ReferenceEntity struct {
+}
+
+// UpdateKeys updates entity referencespdates foreign keys based on an EntityMap.
+func (ent *MinEntity) UpdateKeys(emap *EntityMap) error { return nil }
+
+/////////
+type ErrorEntity struct {
+	loadErrors   []error
+	loadWarnings []error
 }
 
 // AddError adds a loading error to the entity, e.g. from a CSV parse failure
-func (ent *MinEntity) AddError(err error) {
+func (ent *ErrorEntity) AddError(err error) {
 	ent.loadErrors = append(ent.loadErrors, err)
 }
 
 // AddWarning .
-func (ent *MinEntity) AddWarning(err error) {
+func (ent *ErrorEntity) AddWarning(err error) {
 	ent.loadWarnings = append(ent.loadWarnings, err)
 }
 
 // Errors returns validation errors.
-func (ent *MinEntity) Errors() []error { return ent.loadErrors }
+func (ent *ErrorEntity) Errors() []error { return ent.loadErrors }
 
 // Errors returns validation errors.
-func (ent *MinEntity) Warnings() []error { return ent.loadWarnings }
-
-// Filename returns the filename for this entity.
-func (ent *MinEntity) Filename() string { return "" }
-
-// EntityID returns the entity ID.
-func (ent *MinEntity) EntityID() string { return "" }
-
-// UpdateKeys updates entity referencespdates foreign keys based on an EntityMap.
-func (ent *MinEntity) UpdateKeys(emap *EntityMap) error { return nil }
+func (ent *ErrorEntity) Warnings() []error { return ent.loadWarnings }
 
 /////////////
 
@@ -126,6 +161,8 @@ func (ent *Timestamps) UpdateTimestamps() {
 
 type BaseEntity struct {
 	MinEntity
+	ExtraEntity
+	ErrorEntity
 	DatabaseEntity
 	FeedVersionEntity
 	Timestamps

tl/gtfs_shape.go

@@ -63,7 +63,7 @@ func NewShapeFromShapes(shapes []Shape) Shape {
 			ent.AddError(err)
 		}
 		// For tests...
-		if v, ok := shape.Extra()["expect_error"]; len(v) > 0 && ok {
+		if v, ok := shape.GetExtra("expect_error"); len(v) > 0 && ok {
 			ent.SetExtra("expect_error", v)
 		}
 	}

tl/gtfs_stop_time.go

@@ -24,40 +24,15 @@ type StopTime struct {
 	Timepoint         Int
 	Interpolated      Int `csv:"-"` // interpolated times: 0 for provided, 1 interpolated // TODO: 1 for shape, 2 for straight-line
 	FeedVersionID     int `csv:"-"`
-	extra             []string
-	loadErrors        []error
-	loadWarnings      []error
+	ErrorEntity
+	ExtraEntity
 }
 
 // SetFeedVersionID sets the Entity's FeedVersionID.
 func (ent *StopTime) SetFeedVersionID(fvid int) {
 	ent.FeedVersionID = fvid
 }
 
-// AddError adds a loading error to the entity, e.g. from a CSV parse failure
-func (ent *StopTime) AddError(err error) {
-	ent.loadErrors = append(ent.loadErrors, err)
-}
-
-// AddWarning .
-func (ent *StopTime) AddWarning(err error) {
-	ent.loadWarnings = append(ent.loadErrors, err)
-}
-
-// Extra provides any additional fields that were present.
-func (ent *StopTime) Extra() map[string]string {
-	ret := map[string]string{}
-	for i := 0; i < len(ent.extra); i += 2 {
-		ret[ent.extra[i]] = ent.extra[i+1]
-	}
-	return ret
-}
-
-// SetExtra adds a string key, value pair to the entity's extra fields.
-func (ent *StopTime) SetExtra(key string, value string) {
-	ent.extra = append(ent.extra, key, value)
-}
-
 // EntityID returns nothing.
 func (ent *StopTime) EntityID() string {
 	return ""
@@ -67,7 +42,7 @@ func (ent *StopTime) EntityID() string {
 func (ent *StopTime) Errors() []error {
 	// No reflection
 	errs := []error{}
-	errs = append(errs, ent.loadErrors...)
+	errs = append(errs, ent.ErrorEntity.Errors()...)
 	errs = append(errs, tt.CheckPresent("trip_id", ent.TripID)...)
 	errs = append(errs, tt.CheckPresent("stop_id", ent.StopID)...)
 	errs = append(errs, tt.CheckPositiveInt("stop_sequence", ent.StopSequence)...)
@@ -87,11 +62,6 @@ func (ent *StopTime) Errors() []error {
 	return errs
 }
 
-// Warnings for this Entity.
-func (ent *StopTime) Warnings() []error {
-	return ent.loadWarnings
-}
-
 // Filename stop_times.txt
 func (ent *StopTime) Filename() string {
 	return "stop_times.txt"

tl/tests/entity_error_test.go

@@ -19,7 +19,10 @@ func TestEntityErrors(t *testing.T) {
 	}
 	testutil.AllEntities(reader, func(ent tl.Entity) {
 		t.Run(fmt.Sprintf("%s:%s", ent.Filename(), ent.EntityID()), func(t *testing.T) {
-			errs := ent.Errors()
+			var errs []error
+			if extEnt, ok := ent.(tl.EntityWithErrors); ok {
+				errs = extEnt.Errors()
+			}
 			expecterrs := testutil.GetExpectErrors(ent)
 			testutil.CheckErrors(expecterrs, errs, t)
 		})

tlcsv/reader.go

@@ -205,7 +205,7 @@ func (reader *Reader) StopTimesByTripID(tripIDs ...string) chan []tl.StopTime {
 				sid, _ := row.Get("trip_id")
 				if _, ok := set[sid]; ok {
 					ent := tl.StopTime{}
-					loadRowFast(&ent, row)
+					loadRow(&ent, row)
 					m[sid] = append(m[sid], ent)
 				}
 				// If we know the file is grouped, send the stoptimes at transition
@@ -342,7 +342,7 @@ func (reader *Reader) StopTimes() (out chan tl.StopTime) {
 		ent := tl.StopTime{}
 		reader.Adapter.ReadRows(ent.Filename(), func(row Row) {
 			e := tl.StopTime{}
-			loadRowFast(&e, row)
+			loadRow(&e, row)
 			out <- e
 		})
 		close(out)