making the function more robust

Exoutia · Oct 4, 2022 · 69af369 · 69af369
1 parent 4e2c2b5
commit 69af369
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/f1pystats/driver_info.py b/f1pystats/driver_info.py
@@ -7,17 +7,17 @@ def __init__(self, info):
 
     def get_drivers_names(self):
         '''Returns the driver's name list'''
-        return [" ".join([i["givenName"], i["familyName"]]) for i in self.info]
+        return [" ".join([i["givenName"], i["familyName"]]) for i in self.info["MRData"]["DriverTable"]["Drivers"]]
 
     def get_drivers_dob(self):
         '''Returns the driver's date of birth list'''
-        return [i["dateOfBirth"] for i in self.info]
+        return [i["dateOfBirth"] for i in self.info["MRData"]["DriverTable"]["Drivers"]]
 
     def get_drivers_nationality(self):
         """Returns the list of driver's nationality"""
-        return [i["nationality"] for i in self.info]
+        return [i["nationality"] for i in self.info["MRData"]["DriverTable"]["Drivers"]]
 
     def get_drivers_number(self):
         """Returns the list of driver's number"""
-        return [i["permanentNumber"] for i in self.info]
+        return [i["permanentNumber"] for i in self.info["MRData"]["DriverTable"]["Drivers"]]
 
diff --git a/f1pystats/f1pystats.py b/f1pystats/f1pystats.py
@@ -29,7 +29,7 @@ def get_drivers(year: int, round_num: int =None):
     page = requests.get(link, timeout=15)
 
     json_data = page.json()
-    driver_info_json = json_data["MRData"]["DriverTable"]["Drivers"]
+    driver_info_json = json_data
     dr_info = DriverInfo(driver_info_json)
 
     dr_name = dr_info.get_drivers_names()