Add regex to split names properly

youtube_api.py

@@ -30,6 +30,7 @@ def get_playlist_length(youtube, playlist_id):
     return playlist_info['items'][0]['contentDetails']['itemCount']
 
 def main():
+    import re
     parser = argparse.ArgumentParser(description="Add a song to a YouTube playlist.")
     parser.add_argument("channel_name", help="Name of the YouTube channel")
     parser.add_argument("song", help="Name of the song to be added")
@@ -38,9 +39,10 @@ def main():
 
     channel_name = args.channel_name
     song = args.song
-    # this assumes a song is passed as
-    # https://youtu.be/eEnyi9L6KP4
-    song = str(song).split("/")[-1]
+    # this assumes a song is passed as in either of the following ways
+    # https://youtu.be/eEnyi9L6KP4 --> split by "/"
+    # https://youtube.com/watch?v=Ez-gizOF0Wo --> split by "="
+    song = re.split(r"/|=", song)[-1]
 
     if channel_name not in channels:
         print(f"Error: Channel '{channel_name}' not found.")