Skip to content

Instantly share code, notes, and snippets.

@prakhar1989
Created August 14, 2012 14:32
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save prakhar1989/3349803 to your computer and use it in GitHub Desktop.
Save prakhar1989/3349803 to your computer and use it in GitHub Desktop.
Python interface to the Grouplens database
import sqlite3
class User():
def __init__(self, id, age, occupation, male):
self.id = id
self.male = male
self.age = age
self.occupation = occupation
self.movie_ratings = {}
def set_movie_ratings(self):
"""Fills up the movie_ratings dict with movies
and ratings by the current user. Call this right
after making a user"""
conn = sqlite3.connect("data.db")
c = conn.cursor()
t = (self.id,)
c.execute("select * from data where user_id = ?", t)
for r in c.fetchall():
self.movie_ratings[r[1]] = r[2]
c.close()
def movies(self):
"""Returns a list of movies that the user has rated"""
return self.movie_ratings.keys()
def rating(self, movie_name):
"""Returns the rating given by user to the movie.
Returns None if no rating is given"""
if movie_name in self.movie_ratings.keys():
return self.movie_ratings[movie_name]
return None
def get_user_stats(self):
"""Returns the average rating by the user and the
standard deviation in ratings. To be used in
calculating the pearson coefficient"""
conn = sqlite3.connect("data.db")
c = conn.cursor()
t = (self.id, )
c.execute("select * from user_stats where user_id = ?", t)
s = c.fetchone()
c.close
return (s[1], s[2])
def make_user_object(i):
""" Returns a user object with the id as i"""
conn = sqlite3.connect("data.db")
c = conn.cursor()
t = (i, )
c.execute("SELECT * FROM users where user_id = ?", t)
u = c.fetchone()
c.close()
user = User(id = u[0], male = u[2], age = u[1], occupation = u[3])
return user
#USAGE
u = make_user_object(1)
u.set_movie_ratings()
print u.age
print u.rating(24) #rating for the movie with id 24
print u.get_user_stats()[1] #average rating for the user
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment