Coverage for pygeoapi/formatter/csv_.py: 40%

1# ================================================================= 

2# 

3# Authors: Tom Kralidis <tomkralidis@gmail.com> 

4# 

5# Copyright (c) 2022 Tom Kralidis 

6# 

7# Permission is hereby granted, free of charge, to any person 

8# obtaining a copy of this software and associated documentation 

9# files (the "Software"), to deal in the Software without 

10# restriction, including without limitation the rights to use, 

11# copy, modify, merge, publish, distribute, sublicense, and/or sell 

12# copies of the Software, and to permit persons to whom the 

13# Software is furnished to do so, subject to the following 

14# conditions: 

15# 

16# The above copyright notice and this permission notice shall be 

17# included in all copies or substantial portions of the Software. 

18# 

19# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 

20# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES 

21# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 

22# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 

23# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 

24# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 

25# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 

26# OTHER DEALINGS IN THE SOFTWARE. 

27# 

28# ================================================================= 

29 

30import csv 

31import io 

32import logging 

33 

34from pygeoapi.formatter.base import BaseFormatter, FormatterSerializationError 

35 

36LOGGER = logging.getLogger(__name__) 

37 

38 

39class CSVFormatter(BaseFormatter): 

40 """CSV formatter""" 

41 

42 def __init__(self, formatter_def: dict): 

43 """ 

44 Initialize object 

45 

46 :param formatter_def: formatter definition 

47 

48 :returns: `pygeoapi.formatter.csv_.CSVFormatter` 

49 """ 

50 

51 geom = formatter_def.get('geom', False) 

52 

53 super().__init__({'name': 'csv', 'geom': geom}) 

54 self.mimetype = 'text/csv; charset=utf-8' 

55 self.f = 'csv' 

56 self.extension = 'csv' 

57 

58 def write(self, options: dict = {}, data: dict = None) -> str: 

59 """ 

60 Generate data in CSV format 

61 

62 :param options: CSV formatting options 

63 :param data: dict of GeoJSON data 

64 

65 :returns: string representation of format 

66 """ 

67 

68 is_point = False 

69 try: 

70 fields = list(data['features'][0]['properties'].keys()) 

71 except IndexError: 

72 LOGGER.error('no features') 

73 return str() 

74 

75 if self.geom: 

76 LOGGER.debug('Including point geometry') 

77 if data['features'][0]['geometry']['type'] == 'Point': 

78 fields.insert(0, 'x') 

79 fields.insert(1, 'y') 

80 is_point = True 

81 else: 

82 # TODO: implement wkt geometry serialization 

83 LOGGER.debug('not a point geometry, skipping') 

84 

85 LOGGER.debug(f'CSV fields: {fields}') 

86 

87 try: 

88 output = io.StringIO() 

89 writer = csv.DictWriter(output, fields) 

90 writer.writeheader() 

91 

92 for feature in data['features']: 

93 fp = feature['properties'] 

94 if is_point: 

95 fp['x'] = feature['geometry']['coordinates'][0] 

96 fp['y'] = feature['geometry']['coordinates'][1] 

97 LOGGER.debug(fp) 

98 writer.writerow(fp) 

99 except ValueError as err: 

100 LOGGER.error(err) 

101 raise FormatterSerializationError('Error writing CSV output') 

102 

103 return output.getvalue().encode('utf-8') 

104 

105 def __repr__(self): 

106 return f'<CSVFormatter> {self.name}'