diff options
Diffstat (limited to 'ext_test.py')
| -rw-r--r-- | ext_test.py | 13 |
1 files changed, 7 insertions, 6 deletions
diff --git a/ext_test.py b/ext_test.py index 3af637b..9e64223 100644 --- a/ext_test.py +++ b/ext_test.py @@ -61,12 +61,13 @@ class RunOnTest(SimpleExtension): for d in self.test_stream.get_epoch_iterator(as_dict=True): input_values = [d[k.name] for k in self.inputs] output_values = self.function(*input_values) - if 'destination' in self.outputs: - destination = output_values[self.outputs.index('destination')] - dest_outcsv.writerow([d['trip_id'][0], destination[0, 0], destination[0, 1]]) - if 'duration' in self.outputs: - duration = output_values[self.outputs.index('duration')] - time_outcsv.writerow([d['trip_id'][0], int(round(duration[0]))]) + for i in range(d['trip_id'].shape[0]): + if 'destination' in self.outputs: + destination = output_values[self.outputs.index('destination')] + dest_outcsv.writerow([d['trip_id'][i], destination[i, 0], destination[i, 1]]) + if 'duration' in self.outputs: + duration = output_values[self.outputs.index('duration')] + time_outcsv.writerow([d['trip_id'][i], int(round(duration[i]))]) if 'destination' in self.outputs: dest_outfile.close() |