|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | + |
| 3 | +import csv, logging, os, sys, zipfile |
| 4 | +if sys.platform == 'win32': |
| 5 | + csv.field_size_limit(2**31-1) |
| 6 | +else: |
| 7 | + csv.field_size_limit(sys.maxsize) |
| 8 | +try: |
| 9 | + from urllib import urlretrieve |
| 10 | +except ImportError: |
| 11 | + from urllib.request import urlretrieve |
| 12 | +from scipy.spatial import cKDTree as KDTree |
| 13 | + |
| 14 | +# location of geocode data to download |
| 15 | +GEOCODE_URL = 'http://download.geonames.org/export/dump/cities1000.zip' |
| 16 | +GEOCODE_FILENAME = 'cities1000.txt' |
| 17 | +ALTERNATE_NAMES_URL = 'http://download.geonames.org/export/dump/alternateNames.zip' |
| 18 | +ALTERNATE_NAMES_FILENAME = 'alternateNames.txt' |
| 19 | + |
| 20 | + |
| 21 | +def singleton(cls): |
| 22 | + """Singleton pattern to avoid loading class multiple times |
| 23 | + """ |
| 24 | + instances = {} |
| 25 | + |
| 26 | + def getinstance(): |
| 27 | + if cls not in instances: |
| 28 | + instances[cls] = cls() |
| 29 | + return instances[cls] |
| 30 | + return getinstance |
| 31 | + |
| 32 | + |
| 33 | +@singleton |
| 34 | +class GeocodeData: |
| 35 | + |
| 36 | + def __init__(self, geocode_filename='geocode.csv', country_filename='countries.csv'): |
| 37 | + coordinates, self.__locations = self.__extract(rel_path(geocode_filename)) |
| 38 | + self.__tree = KDTree(coordinates) |
| 39 | + self.__load_countries(rel_path(country_filename)) |
| 40 | + |
| 41 | + def __load_countries(self, country_filename): |
| 42 | + """Load a map of country code to name |
| 43 | + """ |
| 44 | + self.__countries = {} |
| 45 | + with open(country_filename, 'r') as handler: |
| 46 | + for code, name in csv.reader(handler): |
| 47 | + self.__countries[code] = name |
| 48 | + |
| 49 | + def query(self, coordinates): |
| 50 | + """Find closest match to this list of coordinates |
| 51 | + """ |
| 52 | + try: |
| 53 | + distances, indices = self.__tree.query(coordinates, k=1) |
| 54 | + except ValueError as e: |
| 55 | + logging.info('Unable to parse coordinates: {}'.format(coordinates)) |
| 56 | + raise e |
| 57 | + else: |
| 58 | + results = [self.__locations[index] for index in indices] |
| 59 | + for result in results: |
| 60 | + result['country'] = self.__countries.get(result['country_code'], '') |
| 61 | + return results |
| 62 | + |
| 63 | + def __download(self): |
| 64 | + """Download geocode file |
| 65 | + """ |
| 66 | + for url in (GEOCODE_URL, ALTERNATE_NAMES_URL): |
| 67 | + local_filename = os.path.abspath(os.path.basename(url)) |
| 68 | + if not os.path.exists(local_filename): |
| 69 | + logging.info('Downloading: {}'.format(url)) |
| 70 | + urlretrieve(url, local_filename) |
| 71 | + return local_filename |
| 72 | + |
| 73 | + def __extract(self, local_filename): |
| 74 | + """Extract geocode data from zip |
| 75 | + """ |
| 76 | + if os.path.exists(local_filename): |
| 77 | + # open compact CSV |
| 78 | + rows = csv.reader(open(local_filename, 'r', encoding='utf-8')) |
| 79 | + else: |
| 80 | + if not os.path.exists(GEOCODE_FILENAME) or not os.path.exists(ALTERNATE_NAMES_FILENAME): |
| 81 | + # remove files to get updated data |
| 82 | + if os.path.exists(GEOCODE_FILENAME): |
| 83 | + os.remove(GEOCODE_FILENAME) |
| 84 | + if os.path.exists(ALTERNATE_NAMES_FILENAME): |
| 85 | + os.remove(ALTERNATE_NAMES_FILENAME) |
| 86 | + |
| 87 | + for url, filename in ((GEOCODE_URL, GEOCODE_FILENAME), (ALTERNATE_NAMES_URL, ALTERNATE_NAMES_FILENAME)): |
| 88 | + downloadedFile = os.path.abspath(os.path.basename(url)) # Re-derive local filename |
| 89 | + if not os.path.exists(downloadedFile): # Check if already downloaded in __download loop? No, this is clearer. |
| 90 | + downloadedFile = self.__download() # Wait, __download downloads BOTH now. |
| 91 | + |
| 92 | + # Actually, let's simplify. __download ensures files are there. |
| 93 | + # Just unzip them. |
| 94 | + pass |
| 95 | + |
| 96 | + # Re-do this logic properly. |
| 97 | + # 1. Download both if missing on extraction needed. |
| 98 | + self.__download() |
| 99 | + |
| 100 | + logging.info('Extracting: {}'.format(GEOCODE_FILENAME)) |
| 101 | + with zipfile.ZipFile(os.path.abspath(os.path.basename(GEOCODE_URL))) as z: |
| 102 | + with open(GEOCODE_FILENAME, 'wb') as fp: |
| 103 | + fp.write(z.read(GEOCODE_FILENAME)) |
| 104 | + |
| 105 | + logging.info('Extracting: {}'.format(ALTERNATE_NAMES_FILENAME)) |
| 106 | + with zipfile.ZipFile(os.path.abspath(os.path.basename(ALTERNATE_NAMES_URL))) as z: |
| 107 | + with open(ALTERNATE_NAMES_FILENAME, 'wb') as fp: |
| 108 | + fp.write(z.read(ALTERNATE_NAMES_FILENAME)) |
| 109 | + |
| 110 | + # Load Arabic names |
| 111 | + arabic_names = {} |
| 112 | + logging.info('Loading Arabic names...') |
| 113 | + with open(ALTERNATE_NAMES_FILENAME, 'r', encoding='utf-8') as f: |
| 114 | + reader = csv.reader(f, delimiter='\t') |
| 115 | + for row in reader: |
| 116 | + # alternateNameId, geonameId, isolanguage, alternateName, isPreferred, isShort, isColloquial, isHistoric |
| 117 | + if len(row) > 3 and row[2] == 'ar': |
| 118 | + geoname_id = row[1] |
| 119 | + arabic_name = row[3] |
| 120 | + # Prefer short names or just overwrite? Let's just take the first one or overwrite. |
| 121 | + # Maybe prefer if isPreferred? row[4] |
| 122 | + arabic_names[geoname_id] = arabic_name |
| 123 | + |
| 124 | + # extract coordinates into more compact CSV for faster loading |
| 125 | + writer = csv.writer(open(local_filename, 'w', encoding='utf-8', newline='')) |
| 126 | + rows = [] |
| 127 | + for row in csv.reader(open(GEOCODE_FILENAME, 'r', encoding='utf-8'), delimiter='\t'): |
| 128 | + geoname_id = row[0] |
| 129 | + latitude, longitude = row[4:6] |
| 130 | + country_code = row[8] |
| 131 | + if latitude and longitude and country_code: |
| 132 | + city = row[1] |
| 133 | + if geoname_id in arabic_names: |
| 134 | + city = arabic_names[geoname_id] |
| 135 | + |
| 136 | + row = latitude, longitude, country_code, city |
| 137 | + writer.writerow(row) |
| 138 | + rows.append(row) |
| 139 | + |
| 140 | + # cleanup downloaded files? Maybe keep them for caching if user wants to rebuild. |
| 141 | + #Original code cleaned up. Let's keep it clean. |
| 142 | + for filename in (os.path.basename(GEOCODE_URL), os.path.basename(ALTERNATE_NAMES_URL), GEOCODE_FILENAME, ALTERNATE_NAMES_FILENAME): |
| 143 | + if os.path.exists(filename): |
| 144 | + os.remove(filename) |
| 145 | + |
| 146 | + # load a list of known coordinates and corresponding __locations |
| 147 | + coordinates, __locations = [], [] |
| 148 | + for latitude, longitude, country_code, city in rows: |
| 149 | + coordinates.append((latitude, longitude)) |
| 150 | + __locations.append(dict(country_code=country_code, city=city)) |
| 151 | + return coordinates, __locations |
| 152 | + |
| 153 | + |
| 154 | +def rel_path(filename): |
| 155 | + """Return the path of this filename relative to the current script |
| 156 | + """ |
| 157 | + return os.path.join(os.getcwd(), os.path.dirname(__file__), filename) |
| 158 | + |
| 159 | + |
| 160 | +def get(coordinate): |
| 161 | + """Search for closest known location to this coordinate |
| 162 | + """ |
| 163 | + gd = GeocodeData() |
| 164 | + return gd.query([coordinate])[0] |
| 165 | + |
| 166 | + |
| 167 | +def search(coordinates): |
| 168 | + """Search for closest known locations to these coordinates |
| 169 | + """ |
| 170 | + gd = GeocodeData() |
| 171 | + return gd.query(coordinates) |
| 172 | + |
| 173 | + |
| 174 | +if __name__ == '__main__': |
| 175 | + # test some coordinate lookups |
| 176 | + city1 = -37.81, 144.96 |
| 177 | + city2 = 31.76, 35.21 |
| 178 | + print(get(city1)) |
| 179 | + print(search([city1, city2])) |
0 commit comments