just small format changes

[pywienerlinien] / gotovienna / routing.py

diff --git a/gotovienna/routing.py b/gotovienna/routing.py
index 279f70a..e70b757 100644 (file)
--- a/gotovienna/routing.py
+++ b/gotovienna/routing.py
@@ -1,12 +1,11 @@
 #!/usr/bin/env python
 # -*- coding: UTF-8 -*-
 
-from BeautifulSoup import BeautifulSoup, NavigableString
+from gotovienna.BeautifulSoup import BeautifulSoup, NavigableString
 from urllib2 import urlopen
 from urllib import urlencode
 from datetime import datetime, time, timedelta
 from textwrap import wrap
-import argparse
 import sys
 import os.path
 import re
@@ -60,6 +59,37 @@ def split_station(station):
     else:
         return (station, 'Wien')
 
+def guess_location_type(location):
+    """Guess type (stop, address, poi) of a location
+
+    >>> guess_location_type('pilgramgasse')
+    'stop'
+
+    >>> guess_location_type('karlsplatz 14')
+    'address'
+
+    >>> guess_location_type('reumannplatz 12/34')
+    'address'
+    """
+    parts = location.split()
+    first_part = parts[0]
+    last_part = parts[-1]
+
+    # Assume all single-word locations are stops
+    if len(parts) == 1:
+        return 'stop'
+
+    # If the last part is numeric, assume address
+    if last_part.isdigit() and len(parts) > 1:
+        return 'address'
+
+    # Addresses with door number (e.g. "12/34")
+    if all(x.isdigit() or x == '/' for x in last_part):
+        return 'address'
+
+    # Sane default - assume it's a stop/station name
+    return 'stop'
+
 def search(origin_tuple, destination_tuple, dtime=None):
     """ build route request
     returns html result (as urllib response)
@@ -74,8 +104,16 @@ def search(origin_tuple, destination_tuple, dtime=None):
     destination, destination_city = split_station(destination)
 
 
-    if not origin_type in POSITION_TYPES or\
-        not destination_type in POSITION_TYPES:
+    if origin_type is None:
+        origin_type = guess_location_type(origin)
+        print 'Guessed origin type:', origin_type
+
+    if destination_type is None:
+        destination_type = guess_location_type(destination)
+        print 'Guessed destination type:', destination_type
+
+    if (origin_type not in POSITION_TYPES or
+            destination_type not in POSITION_TYPES):
         raise ParserError('Invalid position type')
 
     post = defaults.search_post
@@ -116,6 +154,8 @@ class sParser:
 
         return PageType.UNKNOWN
 
+    state = property(check_page)
+
     def get_correction(self):
         names_origin = self.soup.find('select', {'id': 'nameList_origin'})
         names_destination = self.soup.find('select', {'id': 'nameList_destination'})
@@ -123,18 +163,22 @@ class sParser:
         places_destination = self.soup.find('select', {'id': 'placeList_destination'})
         
 
-        if names_origin or names_destination or places_origin or places_destination:
+        if any([names_origin, names_destination, places_origin, places_destination]):
             dict = {}
             
             if names_origin:
-                dict['origin'] = map(lambda x: x.text, names_origin.findAll('option'))
+                dict['origin'] = map(lambda x: x.text, 
+                                     names_origin.findAll('option'))
             if names_destination:
-                dict['destination'] = map(lambda x: x.text, names_destination.findAll('option'))
+                dict['destination'] = map(lambda x: x.text, 
+                                          names_destination.findAll('option'))
                 
             if places_origin:
-                dict['place_origin'] = map(lambda x: x.text, names_origin.findAll('option'))
+                dict['place_origin'] = map(lambda x: x.text, 
+                                           names_origin.findAll('option'))
             if names_destination:
-                dict['place_destination'] = map(lambda x: x.text, names_destination.findAll('option'))
+                dict['place_destination'] = map(lambda x: x.text, 
+                                                names_destination.findAll('option'))
     
             return dict