diff --git a/django_project/gap/providers/observation.py b/django_project/gap/providers/observation.py index 65e0cd49..d9ff0fdb 100644 --- a/django_project/gap/providers/observation.py +++ b/django_project/gap/providers/observation.py @@ -86,7 +86,9 @@ def __init__( self.attributes, key=lambda x: x.attribute.id ) - def _get_data_frame(self, use_separate_time_col=True) -> pd.DataFrame: + def _get_data_frame( + self, use_separate_time_col=True, + use_station_id=False) -> pd.DataFrame: """Create a dataframe from query result. :return: Data frame @@ -122,6 +124,13 @@ def _get_data_frame(self, use_separate_time_col=True) -> pd.DataFrame: }) field_index.append('loc_alt') + # add station id if needed + if use_station_id: + fields.update({ + 's_id': F('station__code') + }) + field_index.append('s_id') + # annotate and select required fields only measurements = self._val.annotate(**fields).values( *(list(fields.keys()) + ['value']) @@ -148,7 +157,7 @@ def _get_data_frame(self, use_separate_time_col=True) -> pd.DataFrame: return df - def _get_headers(self, use_separate_time_col=True): + def _get_headers(self, use_separate_time_col=True, use_station_id=False): """Get list of headers that allign with dataframce columns.""" headers = ['date'] @@ -163,6 +172,10 @@ def _get_headers(self, use_separate_time_col=True): if self.has_altitude_column: headers.append('altitude') + # add station_id + if use_station_id: + headers.append('station_id') + field_indices = [header for header in headers] # add headers @@ -181,13 +194,13 @@ def to_csv_stream(self, suffix='.csv', separator=','): :yield: bytes of csv file :rtype: bytes """ - headers, _ = self._get_headers() + headers, _ = self._get_headers(use_station_id=True) # write headers yield bytes(','.join(headers) + '\n', 'utf-8') # get dataframe - df_pivot = self._get_data_frame() + df_pivot = self._get_data_frame(use_station_id=True) # Write the data in chunks for start in range(0, len(df_pivot), self.csv_chunk_size): @@ -254,12 +267,14 @@ def _to_dict(self) -> dict: has_altitude = self.has_altitude_column output = { 'geometry': json.loads(self.location_input.geometry.json), + 'station_id': '', 'data': [] } # get dataframe df_pivot = self._get_data_frame( - use_separate_time_col=False + use_separate_time_col=False, + use_station_id=True ) for _, row in df_pivot.iterrows(): values = {} @@ -269,9 +284,12 @@ def _to_dict(self) -> dict: 'datetime': row['date'].isoformat(timespec='seconds'), 'values': values }) + if has_altitude: output['altitude'] = row['loc_alt'] + output['station_id'] = row['s_id'] + return output diff --git a/django_project/gap_api/tests/test_measurement_api_airborne.py b/django_project/gap_api/tests/test_measurement_api_airborne.py index a2bfa2f7..3e9181d2 100644 --- a/django_project/gap_api/tests/test_measurement_api_airborne.py +++ b/django_project/gap_api/tests/test_measurement_api_airborne.py @@ -147,6 +147,7 @@ def test_read_point(self): self.assertEqual(len(results), 1) self.assertEqual(results[0]['geometry']['coordinates'], [0, 0]) self.assertEqual(results[0]['altitude'], 1) + self.assertEqual(results[0]['station_id'], 'test-1') self.assertEqual(len(results[0]['data']), 1) self.assertEqual( results[0]['data'][0]['datetime'], '2000-01-01T00:00:00+00:00' @@ -188,8 +189,8 @@ def read_csv(self, response): csv_reader = csv.reader(csv_file) headers = next(csv_reader, None) ordered_headers = [ - 'date', 'time', 'lat', 'lon', 'altitude', 'atmospheric_pressure', - 'temperature' + 'date', 'time', 'lat', 'lon', 'altitude', 'station_id', + 'atmospheric_pressure', 'temperature' ] rows = [] for row in csv_reader: @@ -215,17 +216,17 @@ def test_read_with_bbox(self): headers, rows = self.read_csv(response) self.assertEqual( headers, - ['date', 'time', 'lat', 'lon', 'altitude', 'atmospheric_pressure', - 'temperature'] + ['date', 'time', 'lat', 'lon', 'altitude', 'station_id', + 'atmospheric_pressure', 'temperature'] ) self.assertEqual(len(rows), 2) self.assertEqual( rows[0], - ['2000-02-01', '00:00:00', '10', '10', '2', '200', '2'] + ['2000-02-01', '00:00:00', '10', '10', '2', 'test-1', '200', '2'] ) self.assertEqual( rows[1], - ['2000-03-01', '00:00:00', '100', '100', '3', '300', + ['2000-03-01', '00:00:00', '100', '100', '3', 'test-1', '300', '3'] ) @@ -243,13 +244,13 @@ def test_read_with_bbox(self): headers, rows = self.read_csv(response) self.assertEqual( headers, - ['date', 'time', 'lat', 'lon', 'altitude', 'atmospheric_pressure', - 'temperature'] + ['date', 'time', 'lat', 'lon', 'altitude', 'station_id', + 'atmospheric_pressure', 'temperature'] ) self.assertEqual(len(rows), 1) self.assertEqual( rows[0], - ['2000-02-01', '00:00:00', '10', '10', '2', '200', '2'] + ['2000-02-01', '00:00:00', '10', '10', '2', 'test-1', '200', '2'] ) # Return data with altitude between 1.5-5, should return history 2 & 3 @@ -267,17 +268,18 @@ def test_read_with_bbox(self): headers, rows = self.read_csv(response) self.assertEqual( headers, - ['date', 'time', 'lat', 'lon', 'altitude', 'atmospheric_pressure', - 'temperature'] + ['date', 'time', 'lat', 'lon', 'altitude', 'station_id', + 'atmospheric_pressure', 'temperature'] ) self.assertEqual(len(rows), 2) self.assertEqual( rows[0], - ['2000-02-01', '00:00:00', '10', '10', '2', '200', '2'] + ['2000-02-01', '00:00:00', '10', '10', '2', 'test-1', '200', '2'] ) self.assertEqual( rows[1], - ['2000-03-01', '00:00:00', '100', '100', '3', '300', '3'] + ['2000-03-01', '00:00:00', '100', '100', '3', 'test-1', '300', + '3'] ) def test_read_to_netcdf(self):