d8392a1f4078aa757d1099faf52e96f6ed2348f5
[pywienerlinien] / gotovienna / routing.py
1 #!/usr/bin/env python
2 # -*- coding: UTF-8 -*-
3
4 from BeautifulSoup import BeautifulSoup, NavigableString
5 from urllib2 import urlopen
6 from urllib import urlencode
7 from datetime import datetime, time
8 from textwrap import wrap
9 import argparse
10 import sys
11 import os.path
12
13 from gotovienna import defaults
14
15 POSITION_TYPES = ('stop', 'address', 'poi')
16 TIMEFORMAT = '%H:%M'
17 DEBUGLOG = os.path.expanduser('~/gotoVienna.debug')
18
19 class ParserError(Exception):
20
21     def __init__(self, msg='Parser error'):
22         self.message = msg
23
24 class PageType:
25     UNKNOWN, CORRECTION, RESULT = range(3)
26
27
28 def extract_city(station):
29     """ Extract city from string if present,
30     else return default city
31     
32     >>> extract_city('Karlsplatz, Wien')
33     'Wien'
34     """
35     if len(station.split(',')) > 1:
36         return station.split(',')[-1].strip()
37     else:
38         return 'Wien'
39         
40 def extract_station(station):
41     """ Remove city from string
42     
43     >>> extract_station('Karlsplatz, Wien')
44     'Karlsplatz'
45     """
46     if len(station.split(',')) > 1:
47         return station[:station.rindex(',')].strip()
48     else:
49         return station
50     
51 def split_station(station):
52     """ >>> split_station('Karlsplatz, Wien')
53     ('Karlsplatz', 'Wien')
54     >>> split_station('Karlsplatz')
55     ('Karlsplatz', 'Wien')
56     """
57     print "split:", station
58     if len(station.split(',')) > 1:
59         return (station[:station.rindex(',')].strip(), station.split(',')[-1].strip())
60     else:
61         return (station, 'Wien')
62
63 def search(origin_tuple, destination_tuple, dtime=None):
64     """ build route request
65     returns html result (as urllib response)
66     """
67     if not dtime:
68         dtime = datetime.now()
69
70     origin, origin_type = origin_tuple
71     origin, origin_city = split_station(origin)
72     
73     destination, destination_type = destination_tuple
74     destination, destination_city = split_station(destination)
75
76
77     if not origin_type in POSITION_TYPES or\
78         not destination_type in POSITION_TYPES:
79         raise ParserError('Invalid position type')
80
81     post = defaults.search_post
82     post['name_origin'] = origin
83     post['type_origin'] = origin_type
84     post['name_destination'] = destination
85     post['type_destination'] = destination_type
86     post['itdDateDayMonthYear'] = dtime.strftime('%d.%m.%Y')
87     post['itdTime'] = dtime.strftime('%H:%M')
88     post['place_origin'] = origin_city
89     post['place_destination'] = destination_city
90     params = urlencode(post)
91     url = '%s?%s' % (defaults.action, params)
92
93     try:
94         f = open(DEBUGLOG, 'a')
95         f.write(url + '\n')
96         f.close()
97     except:
98         print 'Unable to write to DEBUGLOG: %s' % DEBUGLOG
99
100     return urlopen(url)
101
102
103 class sParser:
104     """ Parser for search response
105     """
106
107     def __init__(self, html):
108         self.soup = BeautifulSoup(html)
109
110     def check_page(self):
111         if self.soup.find('form', {'id': 'form_efaresults'}):
112             return PageType.RESULT
113
114         if self.soup.find('div', {'class':'form_error'}):
115             return PageType.CORRECTION
116
117         return PageType.UNKNOWN
118
119     def get_correction(self):
120         names_origin = self.soup.find('select', {'id': 'nameList_origin'})
121         names_destination = self.soup.find('select', {'id': 'nameList_destination'})
122         places_origin = self.soup.find('select', {'id': 'placeList_origin'})
123         places_destination = self.soup.find('select', {'id': 'placeList_destination'})
124         
125
126         if names_origin or names_destination or places_origin or places_destination:
127             dict = {}
128             
129             if names_origin:
130                 dict['origin'] = map(lambda x: x.text, names_origin.findAll('option'))
131             if names_destination:
132                 dict['destination'] = map(lambda x: x.text, names_destination.findAll('option'))
133                 
134             if places_origin:
135                 dict['place_origin'] = map(lambda x: x.text, names_origin.findAll('option'))
136             if names_destination:
137                 dict['place_destination'] = map(lambda x: x.text, names_destination.findAll('option'))
138     
139             return dict
140         
141         else:
142             raise ParserError('Unable to parse html')
143
144     def get_result(self):
145         return rParser(str(self.soup))
146
147
148
149 class rParser:
150     """ Parser for routing results
151     """
152
153     def __init__(self, html):
154         self.soup = BeautifulSoup(html)
155         self._overview = None
156         self._details = None
157
158     @classmethod
159     def get_tdtext(cls, x, cl):
160             return x.find('td', {'class': cl}).text
161
162     @classmethod
163     def get_change(cls, x):
164         y = rParser.get_tdtext(x, 'col_change')
165         if y:
166             return int(y)
167         else:
168             return 0
169
170     @classmethod
171     def get_price(cls, x):
172         y = rParser.get_tdtext(x, 'col_price')
173         if y == '*':
174             return 0.0
175         if y.find(','):
176             return float(y.replace(',', '.'))
177         else:
178             return 0.0
179
180     @classmethod
181     def get_date(cls, x):
182         y = rParser.get_tdtext(x, 'col_date')
183         if y:
184             return datetime.strptime(y, '%d.%m.%Y').date()
185         else:
186             return None
187
188     @classmethod
189     def get_time(cls, x):
190         y = rParser.get_tdtext(x, 'col_time')
191         if y:
192             if (y.find("-") > 0):
193                 return map(lambda z: time(*map(int, z.split(':'))), y.split('-'))
194             else:
195                 return map(lambda z: time(*map(int, z.split(':'))), wrap(y, 5))
196         else:
197             return []
198
199     @classmethod
200     def get_duration(cls, x):
201         y = rParser.get_tdtext(x, 'col_duration')
202         if y:
203             return time(*map(int, y.split(":")))
204         else:
205             return None
206
207     def __iter__(self):
208         for detail in self.details():
209             yield detail
210
211     def _parse_details(self):
212         tours = self.soup.findAll('div', {'class': 'data_table tourdetail'})
213
214         trips = map(lambda x: map(lambda y: {
215                         'time': rParser.get_time(y),
216                         'station': map(lambda z: z[2:].strip(),
217                                        filter(lambda x: type(x) == NavigableString, y.find('td', {'class': 'col_station'}).contents)), # filter non NaviStrings
218                         'info': map(lambda x: x.strip(),
219                                     filter(lambda z: type(z) == NavigableString, y.find('td', {'class': 'col_info'}).contents)),
220                     }, x.find('tbody').findAll('tr')),
221                     tours) # all routes
222         return trips
223
224     @property
225     def details(self):
226         """returns list of trip details
227         [ [ { 'time': [datetime.time, datetime.time] if time else [],
228               'station': [u'start', u'end'] if station else [],
229               'info': [u'start station' if station else u'details for walking', u'end station' if station else u'walking duration']
230             }, ... # next trip step
231           ], ... # next trip possibility
232         ]
233         """
234         if not self._details:
235             self._details = self._parse_details()
236
237         return self._details
238
239     def _parse_overview(self):
240
241         # get overview table
242         table = self.soup.find('table', {'id': 'tbl_fahrten'})
243
244         # check if there is an overview table
245         if table and table.findAll('tr'):
246             # get rows
247             rows = table.findAll('tr')[1:] # cut off headline
248
249             overview = map(lambda x: {
250                                'date': rParser.get_date(x),
251                                'time': rParser.get_time(x),
252                                'duration': rParser.get_duration(x), # grab duration
253                                'change': rParser.get_change(x),
254                                'price': rParser.get_price(x),
255                            },
256                            rows)
257         else:
258             raise ParserError('Unable to parse overview')
259
260         return overview
261
262     @property
263     def overview(self):
264         """dict containing
265         date: datetime
266         time: [time, time]
267         duration: time
268         change: int
269         price: float
270         """
271         if not self._overview:
272             try:
273                 self._overview = self._parse_overview()
274             except AttributeError:
275                 f = open(DEBUGLOG, 'w')
276                 f.write(str(self.soup))
277                 f.close()
278
279         return self._overview
280