Add stats

wilty.py

@@ -11,28 +11,50 @@ class Wilty():
         self.save_dir = ".botany"
         self.plant_file = "_plant_data.json"
         self.visit_file = "visitors.json"
+        self.harvest_file = "harvest_file.json"
         self.plants = {}
-        self.loadData()
+        self.stats_model = {
+            "total": 0,
+            "live": 0,
+            "seed": 0,
+            "young": 0,
+            "mature": 0,
+            "succ": 0,
+            "tree": 0,
+            "nonflower": 0,
+            "flower": 0,
+            "gen": 0,
+            "avg_gen": 0,
+            "harv": 0,
+            "avg_harv": 0,
+            }
+        self.border = "-" * 47
 
-    def convertTS(self, ts):
+    def formatTime(self, ts):
-        """Convert a timestamp in seconds to a legible format, e.g. 12d 12h
+        """Format an integer timestamp (seconds) to a timestamp string (days,
-        12m."""
+        hours and minutes), e.g. 12d 12h 12m.
+        """
         days = int(ts / 86400)
         hours = int((ts % 86400) / 3600)
         mins = round(((ts % 86400) % 3600) / 60)
         return str(days) + "d " + str(hours) + "h " + str(mins) + "m"
 
     def getUsers(self):
-        """Load usernames into the plants dictionary as keys."""
+        """Load usernames into the plants dictionary as keys. Only users whose
+        plant data is peer-readable will be added.
+        """
         sys_users = pwd.getpwall()
         for u in sys_users:
-            if os.path.exists("/home/" + u.pw_name + "/" + self.save_dir):
+            sv_dir = "/home/" + u.pw_name + "/" + self.save_dir + "/"
+            if os.path.exists(sv_dir) and \
+                    os.access(sv_dir + u.pw_name + self.plant_file, os.R_OK):
                 self.plants[u.pw_name] = {}
 
     def checkPlant(self, user):
         """Given the username, calculates the last time since the user's
         plant was watered, estimates whether it is dead and updates the plants
-        dictionary."""
+        dictionary.
+        """
         # Get last visitor timestamp and check if later than user's
         visitor_ts = 0
         if self.plants[user]["visitors"] != []:
@@ -40,24 +62,20 @@ class Wilty():
         watered = max(self.plants[user]["last_watered"], visitor_ts)
         time_since = int(time.time()) - watered
         self.plants[user]["time_since"] = time_since
-        # >432000 (5 days) = dead plant
+        # >5 days = dead plant
         if time_since > 432000:
             self.plants[user]["is_dead"] = True
-        self.plants[user]["time_since_fmt"] = self.convertTS(time_since)
+        self.plants[user]["time_since_fmt"] = self.formatTime(time_since)
 
-    def loadData(self):
+    def getRawData(self):
-        """Load plant and visitor data into plants dictionary."""
+        """Load plant data into the plants dictionary."""
         self.getUsers()
         for u in self.plants:
             sv_dir = "/home/" + u + "/" + self.save_dir + "/"
             # Get plant data
-            if os.access(sv_dir + u + self.plant_file, os.R_OK):
             with open(sv_dir + u + self.plant_file) as plant_fh:
                 plant_json = plant_fh.read()
             self.plants[u] = json.loads(plant_json)
-                self.plants[u]["allow_query"] = True
-            else:
-                self.plants[u]["allow_query"] = False
             # Get visitor data
             if os.access(sv_dir + self.visit_file, os.R_OK) and \
                     os.access(sv_dir + self.visit_file, os.W_OK):
@@ -68,37 +86,146 @@ class Wilty():
             else:
                 self.plants[u]["allow_visit"] = False
                 self.plants[u]["visitors"] = []
-            # Update plant watered state if plant data is readable
+            # Get harvest data
-            if self.plants[u]["allow_query"]:
+            if os.access(sv_dir + self.harvest_file, os.R_OK):
+                with open(sv_dir + self.harvest_file) as harvest_fh:
+                    harvest_json = harvest_fh.read()
+                self.plants[u]["harvests"] = json.loads(harvest_json)
+            else:
+                self.plants[u]["harvests"] = {}
+            # Update plant watered state
             self.checkPlant(u)
 
     def listLivePlants(self, *args, **kwargs):
-        """List living plants open to visitors."""
+        """Prints a list of living plants open to visitors. Optionally specify
+        a sort flag to sort the list alphabetically by username or plant type,
+        e.g.  sort=\"plant\". By default, the list is sorted by the longest
+        interval since the last watered time descending.
+        """
         sort_l = []
         for p in self.plants:
-            if self.plants[p]["allow_query"] and \
+            if self.plants[p]["allow_visit"] and not self.plants[p]["is_dead"]:
-                    self.plants[p]["allow_visit"] and \
-                    not self.plants[p]["is_dead"]:
                 sort_l.append((p, self.plants[p]["description"].split()[-1], \
                         self.plants[p]["time_since_fmt"], \
                         self.plants[p]["time_since"]))
         # Sort by column
-        column = kwargs.get("sort", "")
+        if kwargs.get("sort", "") == "user":
-        if column == "user":
             sort_l.sort(key=lambda c: c[0])
-        elif column == "plant":
+        elif kwargs.get("sort", "") == "plant":
             sort_l.sort(key=lambda c: c[1])
         else:
             sort_l.sort(key=lambda c: c[3], reverse=True)
         # Format output
-        border = "-" * 47
+        print(self.border)
-        print(border)
         print("{:<20s}{:<15s}{:<15s}".format("User", "Plant", "Last Watered"))
-        print(border)
+        print(self.border)
         for i in sort_l:
             print("{:<20s}{:<15s}{:<15s}".format(i[0], i[1], i[2]))
-        print(border)
+        print(self.border)
+
+    def countPlantStages(self, key, count):
+        """Count the plants in different stages (seedling, young, mature). Take
+        as arguments a key for plants dictionary lookup and a dictionary to
+        append the results."""
+        if "stage" in self.plants[key]:
+            if self.plants[key]["stage"] == ("mature" or "flowering" or \
+                    "seed-bearing"):
+                count["mature"] += 1
+            elif self.plants[key]["stage"] == "young":
+                count["young"] += 1
+            else:
+                count["seed"] += 1
+        else:
+            count["seed"] += 1
+
+    def countPlantGroups(self, key, count):
+        """Count plants in type groups."""
+        succulents = ["agave", "aloe", "cactus", "lithops", "pachypodium"]
+        trees = ["baobab", "ficus", "palm"]
+        noflowers = ["fern", "moss"]
+        plant_type = self.plants[key]["description"].split()[-1]
+        if plant_type in succulents:
+            count["succ"] += 1
+        elif plant_type in trees:
+            count["tree"] += 1
+        elif plant_type in noflowers:
+            count["nonflower"] += 1
+        elif plant_type != ("seed" or "seedling"):
+            count["flower"] += 1
+
+    def countGenRate(self, key, count):
+        """Count total generation rates."""
+        if "generation" in self.plants[key]:
+            count["gen"] += self.plants[key]["generation"]
+        else:
+            count["gen"] += 1
+
+    def avgGenRate(self, count):
+        """Calculate the average plant generation."""
+        count["avg_gen"] = round(count["gen"] / count["total"], \
+                3)
+        
+    def countHarvests(self, key, count):
+        """Count the total number of harvests."""
+        if self.plants[key]["harvests"] != {}:
+            count["harv"] += len(self.plants[key]["harvests"])
+        else:
+            count["harv"] += 0
+
+    def avgHarvests(self, count):
+        """Calculate the average number of harvests per user."""
+        count["avg_harv"] = round(count["harv"] / count["total"], 3)
+
+    def genStats(self, subset):
+        """Generate garden stats of all plants or only live plants."""
+        stats = dict(self.stats_model)
+        if subset == "live":
+            for p in self.plants:
+                if not self.plants[p]["is_dead"]:
+                    stats["live"] += 1
+                    self.countPlantStages(p, stats)
+                    self.countPlantGroups(p, stats)
+                    self.countGenRate(p, stats)
+                    self.countHarvests(p, stats)
+            stats["total"] = stats["live"]
+        else:
+            for p in self.plants:
+                self.countPlantStages(p, stats)
+                self.countPlantGroups(p, stats)
+                self.countGenRate(p, stats)
+                self.countHarvests(p, stats)
+            stats["total"] = len(self.plants)
+        self.avgGenRate(stats)
+        self.avgHarvests(stats)
+        return stats
+
+    def listStats(self, *args, **kwargs):
+        """Print plant stats. Optionally use a subset flag to get stats on live
+        plants only, e.g. subset=\"live\". By default, show stats based on the
+        total number of plants.
+        """
+        if kwargs.get("subset", "all") == "live":
+            stats = self.genStats("live")
+        else:
+            stats = self.genStats("all")
+        print("Garden Stats (" + kwargs.get("subset", "all") + " plants)")
+        print(self.border)
+        print(
+            "Avg. generation: " +  str(stats["avg_gen"]) +
+            "\nAvg. harvests per user: " +  str(stats["avg_harv"]) +
+            "\nSeedlings: " + str(stats["seed"]) +
+            "\nYoung plants: " + str(stats["young"]) + 
+            "\nMature plants: " + str(stats["mature"]) +  
+            "\nFlowering shrubs: " + str(stats["flower"]) + 
+            "\nNon-flowering plants: " + str(stats["nonflower"]) + 
+            "\nSucculents: " + str(stats["succ"]) +
+            "\nTrees: " + str(stats["tree"]) + 
+            "\nTotal plants: " + str(stats["total"])
+            )
+        print(self.border)
 
 
 instance = Wilty()
+instance.getRawData()
 instance.listLivePlants()
+instance.listStats()