Merge pull request #6 from cecy07/main

[New Feature] - Solves #5

app.py

@@ -250,6 +250,21 @@ def check_login():
             step=1,
             key="num_tickets",
         )
+
+        st.write(
+            "Select the prefered seat type. If both are selected, the app will try to reserve the first available."
+        )
+        sub_col1, sub_col2 = st.columns(2)
+        with sub_col1:
+            general_seat = st.checkbox("General seat", key="general_seat", value=True)
+        with sub_col2:
+            special_seat = st.checkbox("Special seat", key="special_seat", value=True)
+        if general_seat and not special_seat:
+            seat_type = "R"
+        elif not general_seat and special_seat:
+            seat_type = "S"
+        else:
+            seat_type = "B"
     with col2:
         st.subheader("Email notifications settings")
         st.write("Notify you when the train is available.")
@@ -263,8 +278,12 @@ def check_login():
         st.session_state.email_receivers = email_receivers
 
     col1, col2 = st.columns(2)
-    col1.checkbox("Reserve", key="reserve", value=True)
-    col2.checkbox("Notify", key="notify", value=True)
+    with col1:
+        st.write("Check to reserve the ticket.")
+        st.checkbox("Reserve", key="reserve", value=True)
+    with col2:
+        st.write("Check to Notify to email.")
+        st.checkbox("Notify", key="notify", value=True)
 
     st.markdown("---")
 
@@ -304,6 +323,8 @@ def check_login():
                 str(st.session_state.reserve),
                 "--number-of-tickets",
                 str(st.session_state.num_tickets),
+                "--seat-type",
+                seat_type,
             ]
             print(" ".join(command))
             pid = subprocess.Popen(command)  # open in background

reserve.py

@@ -5,7 +5,9 @@
 import time as python_time
 import argparse
 
-from ktrains.korail.korail import Korail
+from ktrains.korail.korail import AdultPassenger, Korail
+from ktrains.korail.korail import ReserveOption as ReserveOptionKorail
+from ktrains.srt.seat_type import SeatType as ReserveOptionSRT
 from ktrains.srt.srt import SRT
 from ktrains.utils import save_to_log
 from ktrains.manage import (
@@ -34,6 +36,7 @@ def get_trains(
     timeout: int = 5,
     email_sender: str = None,
     email_password: str = None,
+    seat_type: str = "B",
 ):
     """
     Main function to run the script. This function will run forever until the number of tickets is reserved or the number of tries is reached.
@@ -100,10 +103,22 @@ def get_trains(
 
             for train in trains:
                 was_available = train_availability[train.train_number]
-
-                if not was_available and train.seat_available():
+                seat = None
+                ReserveOption = ReserveOptionKorail if mode == "korail" else ReserveOptionSRT
+                if seat_type=="R" and train.general_seat_available() :
+                    seat = ReserveOption.GENERAL_ONLY
+                elif seat_type=="S" and train.special_seat_available():
+                    seat = ReserveOption.SPECIAL_ONLY
+                elif seat_type=="B" and train.seat_available():
+                    seat = ReserveOption.GENERAL_FIRST
+                else:
+                    seat = None
+                if not was_available and seat!=None:
                     if reserve:
-                        ktrains.reserve(train)
+                        if mode == "korail":
+                            ktrains.reserve(train,option=seat)
+                        elif mode == "srt":
+                            ktrains.reserve(train, special_seat=seat)
                         manage_reservation(
                             train,
                             email_sender,
@@ -164,6 +179,7 @@ def get_trains(
     parser.add_argument("--timeout", type=int, default=5)
     parser.add_argument("--email-sender", type=str, default=None)
     parser.add_argument("--email-password", type=str, default=None)
+    parser.add_argument("--seat-type", type=str, default="B")
     args = parser.parse_args()
 
     print(args)