Update app.py

app.py

@@ -38,7 +38,7 @@ class Stats:
         self.names = self.loadnames()
 
         # Build index
-        self.vectors, self.data, self.embeddings = self.index()
+        self.vectors, self.data, self.maxyear, self.embeddings = self.index()
 
     def loadcolumns(self):
         """
@@ -122,12 +122,12 @@ class Stats:
         # Build data dictionary
         vectors = {f'{row["yearID"]}{row["playerID"]}': self.transform(row) for _, row in self.stats.iterrows()}
         data = {f'{row["yearID"]}{row["playerID"]}': dict(row) for _, row in self.stats.iterrows()}
+        maxyear = max(row["yearID"] for _, row in self.stats.iterrows())
 
         embeddings = Embeddings({"transform": Stats.transform})
-
         embeddings.index((uid, vectors[uid], None) for uid in vectors)
 
-        return vectors, data, embeddings
+        return vectors, data, maxyear, embeddings
 
     def metrics(self, name):
         """
@@ -155,13 +155,14 @@ class Stats:
 
         return range(1871, datetime.datetime.today().year), 1950, None
 
-    def search(self, name=None, year=None, row=None, limit=10):
+    def search(self, name=None, year=None, window=None, row=None, limit=10):
         """
         Runs an embeddings search. This method takes either a player-year or stats row as input.
 
         Args:
             name: player name to search
             year: year to search
+            window: limit to window recent seasons
             row: row of stats to search
             limit: max results to return
 
@@ -179,13 +180,17 @@ class Stats:
 
         results, ids = [], set()
         if query is not None:
-            for uid, _ in self.embeddings.search(query, limit * 5):
+            candidates = limit * 100 if window else limit * 5
+            for uid, _ in self.embeddings.search(query, candidates):
                 # Only add unique players
                 if uid[4:] not in ids:
                     result = self.data[uid].copy()
-                    result["link"] = f'https://www.baseball-reference.com/players/{result["nameLast"].lower()[0]}/{result["bbrefID"]}.shtml'
-                    results.append(result)
-                    ids.add(uid[4:])
+
+                    # Add first player if this is a player comparison. Limit results to window, if necessary
+                    if (not ids and not row) or not window or result["yearID"] > self.maxyear - window:
+                        result["link"] = f'https://www.baseball-reference.com/players/{result["nameLast"].lower()[0]}/{result["bbrefID"]}.shtml'
+                        results.append(result)
+                        ids.add(uid[4:])
 
                     if len(ids) >= limit:
                         break
@@ -467,6 +472,9 @@ class Application:
         # Player name
         name = self.name(stats.names, params.get("name"))
 
+        # Limit player-year comparisons using this window
+        window = self.window(None, "window")
+
         # Player metrics
         active, best, metrics = stats.metrics(name)
 
@@ -478,7 +486,7 @@ class Application:
             self.chart(category, metrics)
 
         # Run search
-        results = stats.search(name, year)
+        results = stats.search(name, year, window)
 
         # Display results
         self.table(results, ["link", "nameFirst", "nameLast", "teamID"] + stats.columns[1:])
@@ -495,6 +503,10 @@ class Application:
         st.markdown("Find players with similar statistics.")
 
         stats, category = None, self.category("Batting", "searchcategory")
+
+        # Limit player-year comparisons using this window
+        window = self.window(None, "searchwindow")
+
         with st.form("search"):
             if category == "Batting":
                 stats, columns = self.batting, self.batting.columns[:-6]
@@ -507,7 +519,7 @@ class Application:
             submitted = st.form_submit_button("Search")
             if submitted:
                 # Run search
-                results = stats.search(row=inputs.to_dict(orient="records")[0])
+                results = stats.search(window=window, row=inputs.to_dict(orient="records")[0])
 
                 # Display table
                 self.table(results, ["link", "nameFirst", "nameLast", "teamID"] + stats.columns[1:])
@@ -521,15 +533,16 @@ class Application:
         """
 
         # Get parameters
-        params = {x: st.query_params.get(x) for x in ["category", "name", "year"]}
+        params = {x: st.query_params.get(x) for x in ["category", "name", "window", "year"]}
 
         # Sync parameters with session state
-        if all(x in st.session_state for x in ["category", "name", "year"]):
+        if all(x in st.session_state for x in ["category", "name", "window", "year"]):
             # Copy session year if category and name are unchanged
             params["year"] = str(st.session_state["year"]) if all(params.get(x) == st.session_state[x] for x in ["category", "name"]) else None
 
             # Copy category and name from session state
             params["category"] = st.session_state["category"]
+            params["window"] = st.session_state["window"]
             params["name"] = st.session_state["name"]
 
         return params
@@ -555,6 +568,20 @@ class Application:
         # Radio box component
         return st.radio("Stat", categories, index=default, horizontal=True, key=key)
 
+    def window(self, window, key):
+        """
+        Limit results to last N seasons.
+
+        Args:
+            window: limit to window seasons
+            key: widget key
+
+        Returns:
+            window component
+        """
+
+        return st.number_input("Limit to last N seasons", value=window, step=1, min_value=1, max_value=100, key=key)
+
     def name(self, names, name):
         """
         Builds name input widget.
@@ -670,4 +697,4 @@ if __name__ == "__main__":
 
     # Create and run application
     app = create()
-    app.run()
+    app.run()