feat: add support for column-level location data (#15)

This change adds support for the version 2 of the MOJO binary
format specification, which introduces the extra column-level
location information.

austin/format/mojo.py

@@ -93,15 +93,15 @@ def to_austin(self) -> str:
         return ""
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoString(MojoEvent):
     """MOJO string."""
 
     key: int
     value: str
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoStringReference(MojoEvent):
     """MOJO string reference."""
 
@@ -118,7 +118,7 @@ class MojoIdle(MojoEvent):
     pass
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoMetadata(MojoEvent):
     """MOJO metadata."""
 
@@ -130,7 +130,7 @@ def to_austin(self) -> str:
         return f"# {self.key}: {self.value}\n"
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoStack(MojoEvent):
     """MOJO stack."""
 
@@ -142,17 +142,20 @@ def to_austin(self) -> str:
         return f"P{self.pid};T{int(self.tid, 16)}"
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoFrame(MojoEvent):
     """MOJO frame."""
 
     key: int
     filename: MojoStringReference
     scope: MojoStringReference
     line: int
+    line_end: int = 0
+    column: int = 0
+    column_end: int = 0
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoKernelFrame(MojoEvent):
     """MOJO kernel frame."""
 
@@ -163,7 +166,7 @@ def to_austin(self) -> str:
         return f";kernel:{self.scope}:0"
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoSpecialFrame(MojoEvent):
     """MOJO special frame."""
 
@@ -174,7 +177,7 @@ def to_austin(self) -> str:
         return f";:{self.label}:"
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoFrameReference(MojoEvent):
     """MOJO frame reference."""
 
@@ -185,7 +188,7 @@ def to_austin(self) -> str:
         return f";{self.frame.filename.to_austin()}:{self.frame.scope.to_austin()}:{self.frame.line}"
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoMetric(MojoEvent):
     """MOJO metric."""
 
@@ -197,7 +200,7 @@ def to_austin(self) -> str:
         return f" {self.value}\n"
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class MojoFullMetrics(MojoEvent):
     """MOJO full metrics."""
 
@@ -339,7 +342,16 @@ def parse_frame(self) -> t.Generator[MojoFrame, None, None]:
         scope = MojoStringReference(self._lookup_string())
         line = self.read_int()
 
-        frame = MojoFrame(frame_key, filename, scope, line)
+        if self.mojo_version == 1:
+            line_end = column = column_end = 0
+        else:
+            line_end = self.read_int()
+            column = self.read_int()
+            column_end = self.read_int()
+
+        frame = MojoFrame(
+            frame_key, filename, scope, line, line_end, column, column_end
+        )
 
         self._frame_map[self.ref(frame_key)] = frame
 
@@ -421,7 +433,7 @@ def parse_event(self) -> t.Generator[t.Optional[MojoEvent], None, None]:
 
         try:
             for event in t.cast(dict, self.__handlers__)[event_id](self):
-                event.raw = bytes(self._last_bytes)
+                object.__setattr__(event, "raw", bytes(self._last_bytes))
                 self._last_bytes.clear()
                 yield event
         except KeyError as exc:

test/data/column.py

@@ -0,0 +1,20 @@
+from time import sleep
+
+def lazy(n):
+    for _ in range(n):
+        sleep(0.1)
+        yield _
+
+
+def fib(n):
+    a, b = 0, 1
+    for _ in range(n):
+        yield a
+        a, b = b, a + b
+
+a = [
+    list(fib(_))
+    for _ in lazy(30)
+]
+
+print(a)

test/format/test_mojo.py

@@ -29,7 +29,7 @@
 
 import pytest
 
-from austin.format.mojo import MojoFile
+from austin.format.mojo import MojoFile, MojoFrame, MojoString, MojoStringReference
 from austin.format.mojo import main
 from austin.format.mojo import to_varint
 
@@ -58,3 +58,103 @@ def test_mojo_varint():
         buffer.write(b"MOJ\0" + to_varint(n))
         buffer.seek(0)
         assert MojoFile(buffer).read_int() == n
+
+
+def test_mojo_column_info():
+    with (DATA / "column.mojo").open("rb") as stream:
+        frames = {
+            _
+            for _ in MojoFile(stream).parse()
+            if isinstance(_, MojoFrame) and _.filename.string.value == "/tmp/column.py"
+        }
+        assert frames == {
+            MojoFrame(
+                key=1289736945696,
+                filename=MojoStringReference(
+                    string=MojoString(key=20271280, value="/tmp/column.py")
+                ),
+                scope=MojoStringReference(
+                    string=MojoString(key=28930616, value="<module>")
+                ),
+                line=15,
+                line_end=18,
+                column=5,
+                column_end=2,
+            ),
+            MojoFrame(
+                key=1293162643485,
+                filename=MojoStringReference(
+                    string=MojoString(key=20271280, value="/tmp/column.py")
+                ),
+                scope=MojoStringReference(
+                    string=MojoString(key=20364976, value="lazy")
+                ),
+                line=5,
+                line_end=5,
+                column=9,
+                column_end=19,
+            ),
+            MojoFrame(
+                key=1293180469286,
+                filename=MojoStringReference(
+                    string=MojoString(key=20271280, value="/tmp/column.py")
+                ),
+                scope=MojoStringReference(string=MojoString(key=20357744, value="fib")),
+                line=11,
+                line_end=13,
+                column=5,
+                column_end=24,
+            ),
+            MojoFrame(
+                key=1276044640259,
+                filename=MojoStringReference(
+                    string=MojoString(key=20271280, value="/tmp/column.py")
+                ),
+                scope=MojoStringReference(
+                    string=MojoString(key=28930552, value="<listcomp>")
+                ),
+                line=15,
+                line_end=18,
+                column=5,
+                column_end=2,
+            ),
+            MojoFrame(
+                key=1289736945703,
+                filename=MojoStringReference(
+                    string=MojoString(key=20271280, value="/tmp/column.py")
+                ),
+                scope=MojoStringReference(
+                    string=MojoString(key=28930616, value="<module>")
+                ),
+                line=20,
+                line_end=20,
+                column=1,
+                column_end=9,
+            ),
+            MojoFrame(
+                key=1293162643483,
+                filename=MojoStringReference(
+                    string=MojoString(key=20271280, value="/tmp/column.py")
+                ),
+                scope=MojoStringReference(
+                    string=MojoString(key=20364976, value="lazy")
+                ),
+                line=5,
+                line_end=5,
+                column=9,
+                column_end=19,
+            ),
+            MojoFrame(
+                key=1276044640281,
+                filename=MojoStringReference(
+                    string=MojoString(key=20271280, value="/tmp/column.py")
+                ),
+                scope=MojoStringReference(
+                    string=MojoString(key=28930552, value="<listcomp>")
+                ),
+                line=16,
+                line_end=16,
+                column=5,
+                column_end=17,
+            ),
+        }