Develop

stock_prediction.py

@@ -4,10 +4,23 @@
 
 
 # The percentage by which a stock has to beat the S&P500 to be considered a 'buy'
-OUTPERFORMANCE = 10
+#OUTPERFORMANCE = 90
 
+def get_outer_performance():
+    while True:
+        try:
+            x = int(input("Enter min outer performance: "))
+        except ValueError:
+            print("Not an integer!")
+            continue
+        else:
+            if (x < 0 or x >100):
+                print("Valuse must be >=0 and <=100")
+                continue
+            else:
+                return (x)
 
-def build_data_set():
+def build_data_set(outer_performance):
     """
     Reads the keystats.csv file and prepares it for scikit-learn
     :return: X_train and y_train numpy arrays
@@ -22,15 +35,17 @@ def build_data_set():
         status_calc(
             training_data["stock_p_change"],
             training_data["SP500_p_change"],
-            OUTPERFORMANCE,
+            outer_performance,
         )
     )
 
     return X_train, y_train
 
 
 def predict_stocks():
-    X_train, y_train = build_data_set()
+    outer_performance = get_outer_performance()
+
+    X_train, y_train = build_data_set(outer_performance)
     # Remove the random_state parameter to generate actual predictions
     clf = RandomForestClassifier(n_estimators=100, random_state=0)
     clf.fit(X_train, y_train)
@@ -49,7 +64,7 @@ def predict_stocks():
     else:
         invest_list = z[y_pred].tolist()
         print(
-            f"{len(invest_list)} stocks predicted to outperform the S&P500 by more than {OUTPERFORMANCE}%:"
+            f"{len(invest_list)} stocks predicted to outperform the S&P500 by more than {outer_performance}%:"
         )
         print(" ".join(invest_list))
         return invest_list