Merge pull request #478 from 6mxo89sqh/today-tomorrow

Added feature to search for today and tomorrow

src/utils.py

@@ -157,14 +157,18 @@ def get_dose_num(collected_details):
 
     return 1
 
-def start_date_search():
+def start_date_search(find_option):
     # Get search start date
+    print('\nSearch from when?')
     start_date = input(
-        "\nSearch from when?\nUse 1 for today, 2 for tomorrow, or provide a date in the format dd-mm-yyyy. Default 2: "
+        "\nUse 1 for today, 2 for tomorrow, or provide a date in the format dd-mm-yyyy. Default 2: "
+    ) if find_option == 1 else input(
+        "\nUse 1 for today, 2 for tomorrow, 3 for today and tomorrow, or provide a date in the format dd-mm-yyyy. "
+        "Default 2: "
     )
     if not start_date:
         start_date = 2
-    elif start_date in ["1", "2"]:
+    elif start_date in ["1", "2", "3"]:
         start_date = int(start_date)
     else:
         try:
@@ -291,11 +295,11 @@ def collect_user_details(request_header):
                 os.system("pause")
                 sys.exit(1)
         else:
-            start_date=start_date_search()
+            start_date= start_date_search(find_option)
 
     else:
         # Non vaccinated
-        start_date=start_date_search()
+        start_date= start_date_search(find_option)
 
     fee_type = get_fee_type_preference()
 
@@ -430,7 +434,7 @@ def check_by_district(
 
                 if "centers" in resp:
                     print(
-                        f"Centers available in {location['district_name']} from {start_date} as of {today.strftime('%Y-%m-%d %H:%M:%S')}: {len(resp['centers'])}"
+                        f"Centers available in {location['district_name']} {'from' if find_option == 1 else 'for'} {start_date} as of {today.strftime('%Y-%m-%d %H:%M:%S')}: {len(resp['centers'])} "
                     )
                     options += viable_options(
                         resp, minimum_slots, min_age_booking, fee_type, dose_num
@@ -506,7 +510,7 @@ def check_by_pincode(
 
                 if "centers" in resp:
                     print(
-                        f"Centers available in {location['pincode']} from {start_date} as of {today.strftime('%Y-%m-%d %H:%M:%S')}: {len(resp['centers'])}"
+                        f"Centers available in {location['pincode']} {'from' if find_option == 1 else 'for'} {start_date} as of {today.strftime('%Y-%m-%d %H:%M:%S')}: {len(resp['centers'])}"
                     )
                     options += viable_options(
                         resp, minimum_slots, min_age_booking, fee_type, dose_num
@@ -634,75 +638,30 @@ def check_and_book(
         minimum_slots = kwargs["min_slots"]
         refresh_freq = kwargs["ref_freq"]
         # auto_book = kwargs["auto_book"]
-        start_date = kwargs["start_date"]
+        start_dates = []
+        input_start_date = kwargs["start_date"]
         vaccine_type = kwargs["vaccine_type"]
         fee_type = kwargs["fee_type"]
         mobile = kwargs["mobile"]
         # captcha_automation = kwargs['captcha_automation']
         dose_num = kwargs['dose_num']
 
-        if isinstance(start_date, int) and start_date == 2:
-            start_date = (
+        if isinstance(input_start_date, int) and input_start_date in [1, 3]:
+            start_dates.append(datetime.datetime.today().strftime("%d-%m-%Y"))
+        if isinstance(input_start_date, int) and input_start_date in [2, 3]:
+            start_dates.append((
                     datetime.datetime.today() + datetime.timedelta(days=1)
-            ).strftime("%d-%m-%Y")
-        elif isinstance(start_date, int) and start_date == 1:
-            start_date = datetime.datetime.today().strftime("%d-%m-%Y")
-        else:
-            pass
-
-        if search_option == 3:
-            options = check_by_district(
-                find_option,
-                request_header,
-                vaccine_type,
-                location_dtls,
-                start_date,
-                minimum_slots,
-                min_age_booking,
-                fee_type,
-                dose_num,
-                beep_required=False
-            )
-
-            if not isinstance(options, bool):
-                pincode_filtered_options = []
-                for option in options:
-                    for location in pin_code_location_dtls:
-                        if int(location["pincode"]) == int(option["pincode"]):
-                            # ADD this filtered PIN code option
-                            pincode_filtered_options.append(option)
-                            for _ in range(2):
-                                beep(location["alert_freq"], 150)
-                options = pincode_filtered_options
-
-        elif search_option == 2:
-            options = check_by_district(
-                find_option,
-                request_header,
-                vaccine_type,
-                location_dtls,
-                start_date,
-                minimum_slots,
-                min_age_booking,
-                fee_type,
-                dose_num,
-                beep_required=True
-            )
-        else:
-            options = check_by_pincode(
-                find_option,
-                request_header,
-                vaccine_type,
-                location_dtls,
-                start_date,
-                minimum_slots,
-                min_age_booking,
-                fee_type,
-                dose_num
-            )
+            ).strftime("%d-%m-%Y"))
+        if not isinstance(input_start_date, int):
+            start_dates.append(input_start_date)
 
-        if isinstance(options, bool):
-            return False
+        options = []
+        for start_date in start_dates:
+            options_for_date = get_options_for_date(dose_num, fee_type, find_option, location_dtls, min_age_booking, minimum_slots,
+                                       pin_code_location_dtls, request_header, search_option, start_date, vaccine_type)
+            if isinstance(options_for_date, bool):
+                return False
+            options.extend(options_for_date)
 
         options = sorted(
             options,
@@ -827,6 +786,63 @@ def check_and_book(
 
             # tried all slots of all centers but still not able to book then look for current status of centers
             return True
+
+
+def get_options_for_date(dose_num, fee_type, find_option, location_dtls, min_age_booking, minimum_slots,
+                         pin_code_location_dtls, request_header, search_option, start_date, vaccine_type):
+    if search_option == 3:
+        options = check_by_district(
+            find_option,
+            request_header,
+            vaccine_type,
+            location_dtls,
+            start_date,
+            minimum_slots,
+            min_age_booking,
+            fee_type,
+            dose_num,
+            beep_required=False
+        )
+
+        if not isinstance(options, bool):
+            pincode_filtered_options = []
+            for option in options:
+                for location in pin_code_location_dtls:
+                    if int(location["pincode"]) == int(option["pincode"]):
+                        # ADD this filtered PIN code option
+                        pincode_filtered_options.append(option)
+                        for _ in range(2):
+                            beep(location["alert_freq"], 150)
+            options = pincode_filtered_options
+
+    elif search_option == 2:
+        options = check_by_district(
+            find_option,
+            request_header,
+            vaccine_type,
+            location_dtls,
+            start_date,
+            minimum_slots,
+            min_age_booking,
+            fee_type,
+            dose_num,
+            beep_required=True
+        )
+    else:
+        options = check_by_pincode(
+            find_option,
+            request_header,
+            vaccine_type,
+            location_dtls,
+            start_date,
+            minimum_slots,
+            min_age_booking,
+            fee_type,
+            dose_num
+        )
+    return options
+
+
 def get_vaccine_preference():
     print(
         "It seems you're trying to find a slot for your first dose. Do you have a vaccine preference?"