POST on same page when removing from playlist, so url doesn't change

1 files changed, 10 insertions, 8 deletions

diff --git a/youtube/youtube.py b/youtube/youtube.py
index a9bd3c4..93b947f 100644
--- a/youtube/youtube.py
+++ b/youtube/youtube.py
@@ -53,24 +53,26 @@ def youtube(env, start_response):
             return b'404 Not Found'
 
     elif method == "POST":
+        fields = urllib.parse.parse_qs(env['wsgi.input'].read().decode())
         if path == "/edit_playlist":
-            fields = urllib.parse.parse_qs(env['wsgi.input'].read().decode())
             if fields['action'][0] == 'add':
                 local_playlist.add_to_playlist(fields['playlist_name'][0], fields['video_info_list'])
                 start_response('204 No Content', ())
-            elif fields['action'][0] == 'remove':
-                try:
-                    playlist_name = fields['playlist_page'][0]
-                except KeyError:
-                    playlist_name = fields['playlist_name'][0]
+            else:
+                start_response('400 Bad Request', ())
+                return b'400 Bad Request'
+
+        elif path.startswith("/playlists"):
+            if fields['action'][0] == 'remove':
+                playlist_name = path[11:]
                 local_playlist.remove_from_playlist(playlist_name, fields['video_info_list'])
                 start_response('200 OK', ())
                 return local_playlist.get_playlist_page(playlist_name).encode() 
+                start_response('200 OK',  (('Content-type','text/html'),) )
+                return local_playlist.get_playlist_page(path[10:], query_string=query_string).encode()
             else:
                 start_response('400 Bad Request', ())
                 return b'400 Bad Request'
-
-            
         else:
             start_response('404 Not Found', ())
             return b'404 Not Found'