Update xgboostclassifier.py

machine_learning/xgboostclassifier.py

@@ -1,21 +1,19 @@
 # XGBoost Classifier Example
-from doctest import testmod
-
 from matplotlib import pyplot as plt
 from sklearn.datasets import load_iris
 from sklearn.metrics import plot_confusion_matrix
 from sklearn.model_selection import train_test_split
 from xgboost import XGBClassifier
 
 
-def data_handling(data) -> list:
+def data_handling(data: list) -> tuple:
     # Split dataset into train and test data
     x = data["data"]  # features
     y = data["target"]
     return x, y
 
 
-def xgboost(features, target):
+def xgboost(features: list, target: list):
     classifier = XGBClassifier()
     classifier.fit(features, target)
     return classifier
@@ -54,4 +52,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    testmod(name="main", verbose=True)
+    main()