1 files changed, 9 insertions, 17 deletions
diff --git a/pathod/language/websockets.py b/pathod/language/websockets.py
index b4faf59b..cc00bcf1 100644
--- a/pathod/language/websockets.py
+++ b/pathod/language/websockets.py
@@ -1,10 +1,12 @@
 import random
 import string
+import typing  # noqa
+
+import pyparsing as pp
+
 import mitmproxy.net.websockets
 from mitmproxy.utils import strutils
-import pyparsing as pp
 from . import base, generators, actions, message
-import typing  # noqa
 
 NESTED_LEADER = b"pathod!"
 
@@ -74,7 +76,7 @@ class Times(base.Integer):
     preamble = "x"
 
 
-COMPONENTS = (
+COMPONENTS = [
     OpCode,
     Length,
     # Bit flags
@@ -89,14 +91,13 @@ COMPONENTS = (
     KeyNone,
     Key,
     Times,
-
     Body,
     RawBody,
-)
+]
 
 
 class WebsocketFrame(message.Message):
-    components = COMPONENTS
+    components = COMPONENTS  # type: typing.List[typing.Type[base._Component]]
     logattrs = ["body"]
     # Used for nested frames
     unique_name = "body"
@@ -235,19 +236,10 @@ class WebsocketFrame(message.Message):
         return ":".join([i.spec() for i in self.tokens])
 
 
-class NestedFrame(base.NestedMessage):
+class NestedFrame(message.NestedMessage):
     preamble = "f"
     nest_type = WebsocketFrame
 
 
-COMP = typing.Tuple[
-    typing.Type[OpCode], typing.Type[Length], typing.Type[Fin], typing.Type[RSV1], typing.Type[RSV2], typing.Type[RSV3], typing.Type[Mask],
-    typing.Type[actions.PauseAt], typing.Type[actions.DisconnectAt], typing.Type[actions.InjectAt], typing.Type[KeyNone], typing.Type[Key],
-    typing.Type[Times], typing.Type[Body], typing.Type[RawBody]
-]
-
-
 class WebsocketClientFrame(WebsocketFrame):
-    components = typing.cast(COMP, COMPONENTS + (
-        NestedFrame,
-    ))
+    components = COMPONENTS + [NestedFrame]