Small refactoring.

Tássia Camões Araújo
1 parent 39001e40
Showing 2 changed files with 12 additions and 9 deletions Show diff stats
src/recommender.py
src/strategy.py
@@ -33,6 +33,7 @@ class RecommendationResult:
         Set initial parameters.
         """
         self.item_score = item_score
+        self.size = len(item_score)
     def __str__(self):
         """
@@ -44,14 +45,14 @@ class RecommendationResult:
             str += "%2d: %s\n" % (i,result[i][0])
         return str
-    def get_prediction(self,size=20):
+    def get_prediction(self,limit=20):
         """
         Return prediction based on recommendation size (number of items).
         """
-        if size > len(self.item_score): size = len(self.item_score)
+        if limit > self.size: limit = self.size
         sorted_result = sorted(self.item_score.items(),
                                key=operator.itemgetter(1))
-        return list(reversed(sorted_result[-size:]))
+        return list(reversed(sorted_result[-limit:]))
 class Recommender:
     """
@@ -83,8 +84,8 @@ class Recommender:
         if strategy_str == "col":
             self.strategy = strategy.CollaborativeStrategy(20)
-    def get_recommendation(self,user,limit=20):
+    def get_recommendation(self,user,result_size=20):
         """
         Produces recommendation using previously loaded strategy.
         """
-        return self.strategy.run(self,user,limit)
+        return self.strategy.run(self,user,result_size)
@@ -95,15 +95,17 @@ class ContentBasedStrategy(RecommendationStrategy):
     """
     Content-based recommendation strategy based on Apt-xapian-index.
     """
-    def __init__(self,content):
+    def __init__(self,content,profile_size=50):
         self.description = "Content-based"
         self.content = content
+        self.profile_size = profile_size
     def run(self,rec,user,limit):
         """
         Perform recommendation strategy.
         """
-        profile = user.profile(rec.items_repository,self.content,50)
+        profile = user.profile(rec.items_repository,self.content,
+                               self.profile_size)
         # prepair index for querying user profile
         query = xapian.Query(xapian.Query.OP_OR,profile)
         enquire = xapian.Enquire(rec.items_repository)
@@ -129,7 +131,7 @@ class CollaborativeStrategy(RecommendationStrategy):
         self.clustering = clustering
         self.neighbours = k
-    def run(self,rec,user,limit):
+    def run(self,rec,user,result_size):
         """
         Perform recommendation strategy.
         """
@@ -153,7 +155,7 @@ class CollaborativeStrategy(RecommendationStrategy):
             rset.add_document(m.document.get_docid())
             logging.debug(m.document.get_data())
         # retrieve most relevant packages
-        eset = enquire.get_eset(limit,rset,PkgExpandDecider())
+        eset = enquire.get_eset(result_size,rset,PkgExpandDecider())
         # compose result dictionary
         item_score = {}
         for package in eset: