social_users.py

social_users.py

#! /usr/bin/env python

import pickle
import re

from db import Database

# instanciate the database class
db = Database(
    db_host = '*******',
    db_name = '*******',
    db_username = '****',
    db_password = '***********',
)

# open the pickled country/city mapping and create a country set
populated_city_pkl_file = open('country_populated_city.pkl', 'r')
populated_city_mapping = pickle.load(populated_city_pkl_file)
populated_city_pkl_file.close()

all_city_pkl_file = open('country_city.pkl', 'r')
all_city_mapping = pickle.load(all_city_pkl_file)
country_list = [[k, v['name'].lower()] for k, v in all_city_mapping.items()]
all_city_pkl_file.close()

# instanciate a cache list for faster searching for already determined location strings
# NB: includes a few hacks
cached_locations = [
    ['US', 'USA'],
    ['GB', 'UK'],
    ['US', 'US of A'],
    ['US', 'States'],
    ['US', 'Florida'],
    ['US', 'California'],
]
cached_country_ids = [
    [250, None]
]

def group_by(results, key='location'):
    
    # remove bad chars and sort the results
    results = [r for r in results if not r.update({key: re.sub(r"[\n|\r|\t]", " ", r[key]).strip()})]
    sorted_results = [r for r in sorted(results, key=lambda k: k[key])]
    
    prev_location = None
    grouped_results = []
    unique_counter = -1
    for result in sorted_results:
        if prev_location == result[key]:
            grouped_results[unique_counter].append(result)
        else:
            grouped_results.append([result])
            unique_counter += 1
        prev_location = result[key]
    
    return grouped_results

def parse_location(location):
    
    # return blank to be processed as global if no location is provided
    if not location:
        return []
    
    # split the location by puncutuated phrases
    location_split = [phrase.strip().lower() for phrase in re.split(r'[,|\-|)|(|.|\||#|\:|\\|\/]', location) if phrase]
    
    possible_matches =[]
    city_name = None
    
    # test if the country name or code is already present in the location string
    for country in country_list:
        if country[1] in location_split:
            return [country[0]]
        elif country[0] in location_split:
            return [country[0]]
    
    # check if the city name is in the cache list
    for location in cached_locations:
        if location[1].lower() in location_split:
            return [location[0]]
    
    # iterate through the short list mapping to find a matching city name
    for country_code, country_dict in populated_city_mapping.items():
        for city in country_dict['city_info']:
            if city[0].lower() in location_split:
                possible_matches.append([country_code, city[1]])
                city_name = city[0]
    
    # iterate through the full mapping to find a matching city name
    if not city_name:
        for country_code, country_dict in all_city_mapping.items():
            for city in country_dict['city_info']:
                if city[0].lower() in location_split:
                    possible_matches.append([country_code, city[1]])
                    city_name = city[0]
    
    # sort the matched country list by cities population to increase accuracy by probability
    sorted_matches = [i[0] for i in sorted(possible_matches, key=lambda k: k[1], reverse=True)]
    
    # add the city to the cached list
    if city_name:
        if [sorted_matches[0], city_name] not in cached_locations:
            cached_locations.append([sorted_matches[0], city_name])
        
    return sorted_matches
    
def get_country_id(country_code):
    for c in cached_country_ids:
        if country_code == c[1]:
            return c[0]
    
    sql = "SELECT id FROM countries WHERE abbr2='%s'" % country_code
    country = db.get(sql)
    if country:
        cached_country_ids.append([country['id'], country_code])
        return country['id']
    
    return 250

def get_twitter_users(limit=None):
    sql = "SELECT id, username, name, location FROM twitter_users"
    if limit:
        sql += " LIMIT %s, %s" % (limit[0], limit[1])
    
    return group_by(db.list(sql))
    
def twitter_user_exists(grouped_twitter_users):
    for twitter_user in grouped_twitter_users:
        sql = "SELECT id FROM social_users WHERE twitterId = %s" % twitter_user['id']
        if not db.list(sql):
            return False
    return True
    
def process_twitter_users(data):
    country_id = get_country_id(data['country_code'])
    try:
        db.execute(
            "INSERT INTO social_users (username, full_name, countryId, twitterId) VALUES (%s, %s, %s, %s)", (
                data['username'],
                data['name'],
                country_id,
                data['id'],
            )
        )
        return True
    except:
        return False

def get_facebook_users():
    sql = "SELECT id, username, name, location, gender FROM facebook_users ORDER BY location"
    return db.list(sql)

def process_facebook_users():
    try:
        db.execute(
            "INSERT INTO social_users (username, full_name, gender, countryId, facebookId) VALUES (%s, %s, %s, %s, %s)", (
                data['username'],
                data['name'],
                data['gender'],
                get_country_id(data['country_code']),
                data['id'],
            )
        )
        return True
    except:
        return False

if __name__ == '__main__':
    count = 1
    for grouped_twitter_users in get_twitter_users():
        if not twitter_user_exists(grouped_twitter_users):
            try:
                country_code = parse_location(grouped_twitter_users[0]['location'])[0]
            except IndexError:
                country_code = None
                
            for twitter_user in grouped_twitter_users:
                twitter_user['country_code'] = country_code
                
                if process_twitter_users(twitter_user):
                    print "INSERT(%s): %s" % (count, twitter_user)
                else:
                    print "FAILED(%s): %s" % (count, twitter_user)
                count += 1