6 files changed, 71 insertions, 68 deletions
diff --git a/pathod/language/actions.py b/pathod/language/actions.py
index fc57a18b..3e48f40d 100644
--- a/pathod/language/actions.py
+++ b/pathod/language/actions.py
@@ -50,7 +50,7 @@ class _Action(base.Token):
 
 
 class PauseAt(_Action):
-    unique_name = None  # type: ignore
+    unique_name = None
 
     def __init__(self, offset, seconds):
         _Action.__init__(self, offset)
diff --git a/pathod/language/base.py b/pathod/language/base.py
index c8892748..97871e7e 100644
--- a/pathod/language/base.py
+++ b/pathod/language/base.py
@@ -6,7 +6,8 @@ import pyparsing as pp
 from mitmproxy.utils import strutils
 from mitmproxy.utils import human
 import typing  # noqa
-from . import generators, exceptions
+from . import generators
+from . import exceptions
 
 
 class Settings:
@@ -375,7 +376,7 @@ class OptionsOrValue(_Component):
 
 
 class Integer(_Component):
-    bounds = (None, None)  # type: typing.Tuple[typing.Union[int, None], typing.Union[int , None]]
+    bounds = (None, None)  # type: typing.Tuple[typing.Optional[int], typing.Optional[int]]
     preamble = ""
 
     def __init__(self, value):
@@ -537,43 +538,3 @@ class IntField(_Component):
 
     def spec(self):
         return "%s%s" % (self.preamble, self.origvalue)
-
-
-class NestedMessage(Token):
-
-    """
-        A nested message, as an escaped string with a preamble.
-    """
-    preamble = ""
-    nest_type = None  # type: ignore
-
-    def __init__(self, value):
-        Token.__init__(self)
-        self.value = value
-        try:
-            self.parsed = self.nest_type(
-                self.nest_type.expr().parseString(
-                    value.val.decode(),
-                    parseAll=True
-                )
-            )
-        except pp.ParseException as v:
-            raise exceptions.ParseException(v.msg, v.line, v.col)
-
-    @classmethod
-    def expr(cls):
-        e = pp.Literal(cls.preamble).suppress()
-        e = e + TokValueLiteral.expr()
-        return e.setParseAction(lambda x: cls(*x))
-
-    def values(self, settings):
-        return [
-            self.value.get_generator(settings),
-        ]
-
-    def spec(self):
-        return "%s%s" % (self.preamble, self.value.spec())
-
-    def freeze(self, settings):
-        f = self.parsed.freeze(settings).spec()
-        return self.__class__(TokValueLiteral(strutils.bytes_to_escaped_str(f.encode(), escape_single_quotes=True)))
diff --git a/pathod/language/http.py b/pathod/language/http.py
index 5cd717a9..5a962145 100644
--- a/pathod/language/http.py
+++ b/pathod/language/http.py
@@ -54,7 +54,9 @@ class Method(base.OptionsOrValue):
 
 
 class _HeaderMixin:
-    unique_name = None  # type: ignore
+    @property
+    def unique_name(self):
+        return None
 
     def format_header(self, key, value):
         return [key, b": ", value, b"\r\n"]
@@ -251,7 +253,7 @@ class Response(_HTTPMessage):
         return ":".join([i.spec() for i in self.tokens])
 
 
-class NestedResponse(base.NestedMessage):
+class NestedResponse(message.NestedMessage):
     preamble = "s"
     nest_type = Response
 
diff --git a/pathod/language/http2.py b/pathod/language/http2.py
index 47d6e370..5b27d5bf 100644
--- a/pathod/language/http2.py
+++ b/pathod/language/http2.py
@@ -1,9 +1,9 @@
 import pyparsing as pp
+
 from mitmproxy.net import http
 from mitmproxy.net.http import user_agents, Headers
 from . import base, message
 
-
 """
     Normal HTTP requests:
         <method>:<path>:<header>:<body>
@@ -41,7 +41,9 @@ def get_header(val, headers):
 
 
 class _HeaderMixin:
-    unique_name = None  # type: ignore
+    @property
+    def unique_name(self):
+        return None
 
     def values(self, settings):
         return (
@@ -146,7 +148,7 @@ class Times(base.Integer):
 
 
 class Response(_HTTP2Message):
-    unique_name = None  # type: ignore
+    unique_name = None
     comps = (
         Header,
         Body,
@@ -203,7 +205,7 @@ class Response(_HTTP2Message):
         return ":".join([i.spec() for i in self.tokens])
 
 
-class NestedResponse(base.NestedMessage):
+class NestedResponse(message.NestedMessage):
     preamble = "s"
     nest_type = Response
 
diff --git a/pathod/language/message.py b/pathod/language/message.py
index 6b4c5021..5dda654b 100644
--- a/pathod/language/message.py
+++ b/pathod/language/message.py
@@ -1,8 +1,11 @@
 import abc
-from . import actions, exceptions
-from mitmproxy.utils import strutils
 import typing  # noqa
 
+import pyparsing as pp
+
+from mitmproxy.utils import strutils
+from . import actions, exceptions, base
+
 LOG_TRUNCATE = 1024
 
 
@@ -96,3 +99,46 @@ class Message:
 
     def __repr__(self):
         return self.spec()
+
+
+class NestedMessage(base.Token):
+    """
+        A nested message, as an escaped string with a preamble.
+    """
+    preamble = ""
+    nest_type = None  # type: typing.Optional[typing.Type[Message]]
+
+    def __init__(self, value):
+        super().__init__()
+        self.value = value
+        try:
+            self.parsed = self.nest_type(
+                self.nest_type.expr().parseString(
+                    value.val.decode(),
+                    parseAll=True
+                )
+            )
+        except pp.ParseException as v:
+            raise exceptions.ParseException(v.msg, v.line, v.col)
+
+    @classmethod
+    def expr(cls):
+        e = pp.Literal(cls.preamble).suppress()
+        e = e + base.TokValueLiteral.expr()
+        return e.setParseAction(lambda x: cls(*x))
+
+    def values(self, settings):
+        return [
+            self.value.get_generator(settings),
+        ]
+
+    def spec(self):
+        return "%s%s" % (self.preamble, self.value.spec())
+
+    def freeze(self, settings):
+        f = self.parsed.freeze(settings).spec()
+        return self.__class__(
+            base.TokValueLiteral(
+                strutils.bytes_to_escaped_str(f.encode(), escape_single_quotes=True)
+            )
+        )
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]