flamegraph-textual 0.1.0__py3-none-any.whl

flamegraph_textual/__init__.py

@@ -0,0 +1,25 @@
+__version__ = "0.1.0"
+
+from flamegraph_textual.models import Frame, Profile, SampleType
+from flamegraph_textual.parsers import parse
+from flamegraph_textual.render.app import FlameGraphScroll, FlameshowApp
+from flamegraph_textual.render.flamegraph import FlameGraph, FrameMap, add_array
+from flamegraph_textual.view import FlameGraphView
+
+
+class FlameGraphApp(FlameshowApp):
+    """Convenience alias for embedding a full flamegraph app."""
+
+
+__all__ = [
+    "FlameGraph",
+    "FlameGraphApp",
+    "FlameGraphScroll",
+    "FlameGraphView",
+    "Frame",
+    "FrameMap",
+    "Profile",
+    "SampleType",
+    "add_array",
+    "parse",
+]

flamegraph_textual/colors.py

@@ -0,0 +1,67 @@
+import random
+import logging
+
+from textual.color import Color
+
+logger = logging.getLogger(__name__)
+
+
+class ColorPlatteBase:
+    def __init__(self):
+        self.assigned_color = {}
+
+    def get_color(self, key):
+        if key not in self.assigned_color:
+            self.assigned_color[key] = self.assign_color(key)
+        return self.assigned_color[key]
+
+    def assign_color(self, key):
+        raise NotImplementedError
+
+
+class LinaerColorPlatte(ColorPlatteBase):
+    def __init__(
+        self,
+        start_color=Color.parse("#CD0000"),
+        end_color=Color.parse("#FFE637"),
+    ) -> None:
+        super().__init__()
+        self.assigned_color = {}
+        self.start_color = start_color
+        self.end_color = end_color
+        self.index = 0
+        self.platte = self.generate_platte()
+
+    def assign_color(self, key):
+        color = self.platte[self.index]
+        self.index += 1
+        if self.index == len(self.platte):
+            self.index = 0
+
+        logger.debug("assign color=%s", color)
+        return color
+
+    def generate_platte(self):
+        color_platte = []
+        for factor in range(0, 100, 5):
+            color_platte.append(
+                self.start_color.blend(self.end_color, factor / 100)
+            )
+        return color_platte
+
+
+class FlameGraphRandomColorPlatte(ColorPlatteBase):
+    def __init__(self) -> None:
+        super().__init__()
+        self.assigned_color = {}
+
+    def assign_color(self, *args):
+        return Color(
+            205 + int(50 * random.random()),
+            0 + int(230 * random.random()),
+            0 + int(55 * random.random()),
+        )
+
+
+flamegraph_random_color_platte = FlameGraphRandomColorPlatte()
+linaer_color_platte = LinaerColorPlatte()

flamegraph_textual/const.py

@@ -0,0 +1,4 @@
+VIEW_INFO_COLOR = "#ffffff"
+VIEW_INFO_OTHER_COLOR = "#8884FF"
+SELECTED_PARENTS_BG_COLOR_BLEND_TO = "#8b0000"
+SELECTED_PARENTS_BG_COLOR_BLEND_FACTOR = 0.5

flamegraph_textual/exceptions.py

@@ -0,0 +1,14 @@
+class FlamegraphTextualException(Exception):
+    """Base exception for flamegraph_textual."""
+
+
+class ProfileParseException(FlamegraphTextualException):
+    """Can not parse the profile."""
+
+
+class UsageError(FlamegraphTextualException):
+    """Usage Error."""
+
+
+class RenderException(FlamegraphTextualException):
+    """Error during flamegraph render."""

flamegraph_textual/models.py

@@ -0,0 +1,204 @@
+from dataclasses import dataclass, field
+import datetime
+import logging
+import time
+from typing import Dict, List, Set
+from typing_extensions import Self
+
+from rich.style import Style
+from rich.text import Text
+
+from flamegraph_textual.utils import sizeof
+
+from flamegraph_textual.runtime import r
+
+
+logger = logging.getLogger(__name__)
+
+
+class Frame:
+    def __init__(
+        self, name, _id, children=None, parent=None, values=None, root=None
+    ) -> None:
+        self.name = name
+        self._id = _id
+        if children:
+            self.children = children
+        else:
+            self.children = []
+        self.parent = parent
+        if not values:
+            self.values = []
+        else:
+            self.values = values
+
+        self.root = root
+
+    def pile_up(self, childstack: Self):
+        childstack.parent = self
+
+        for exist_child in self.children:
+            # added to exist, no need to create one
+            if exist_child.name == childstack.name:
+                # some cases, childstack.children total value not equal to
+                # childstack.values
+                # so, we need to add values of "parent" instead of add values
+                # by every child
+                exist_child.values = list(
+                    map(sum, zip(exist_child.values, childstack.values))
+                )
+
+                for new_child in childstack.children:
+                    exist_child.pile_up(new_child)
+                return
+
+        self.children.append(childstack)
+
+    def __eq__(self, other):
+        if isinstance(other, Frame):
+            return self._id == other._id
+        return False
+
+    @property
+    def display_color(self):
+        return r.get_color(self.color_key)
+
+    def humanize(self, sample_unit, value):
+        display_value = value
+        if sample_unit == "bytes":
+            display_value = sizeof(value)
+
+        return display_value
+
+    def __repr__(self) -> str:
+        return f"<Frame #{self._id} {self.name}>"
+
+    def render_detail(self, sample_index: int, sample_unit: str):
+        """
+        render stacked information
+        """
+        detail = []
+        frame = self
+        while frame:
+            lines = self.render_one_frame_detail(
+                frame, sample_index, sample_unit
+            )
+            for line in lines:
+                detail.append(
+                    Text.assemble(
+                        (" ", Style(bgcolor=frame.display_color.rich_color)),
+                        " ",
+                        line,
+                    )
+                )
+            frame = frame.parent
+
+        return Text.assemble(*detail)
+
+    def render_one_frame_detail(
+        self, frame, sample_index: int, sample_unit: str
+    ):
+        raise NotImplementedError
+
+    @property
+    def title(self) -> Text:
+        """Full name which will be displayed in the frame detail panel"""
+        return Text(self.name)
+
+    @property
+    def color_key(self):
+        """Same key will get the same color"""
+        return self.name
+
+    @property
+    def display_name(self):
+        """The name display on the flamegraph"""
+        return self.name
+
+
+@dataclass
+class SampleType:
+    sample_type: str = ""
+    sample_unit: str = ""
+
+
+@dataclass
+class Profile:
+    # required
+    filename: str
+    root_stack: Frame
+    highest_lines: int
+    # total samples is one top most sample, it's a list that contains all
+    # its parents all the way up
+    total_sample: int
+    sample_types: List[SampleType]
+    # int id mapping to Frame
+    id_store: Dict[int, Frame]
+
+    # optional
+    default_sample_type_index: int = -1
+    period_type: SampleType | None = None
+    period: int = 0
+    created_at: datetime.datetime | None = None
+
+    # init by post_init
+    lines: List = field(init=False)
+
+    frameid_to_lineno: Dict[int, int] = field(init=False)
+
+    # Frame grouped by same name
+    name_aggr: Dict[str, List[Frame]] = field(init=False)
+
+    def __post_init__(self):
+        """
+        init_lines must be called before render
+        """
+        t1 = time.time()
+        logger.info("start to create lines...")
+
+        root = self.root_stack
+
+        lines = [
+            [root],
+        ]
+        frameid_to_lineno = {0: 0}
+        current = root.children
+        line_no = 1
+
+        while len(current) > 0:
+            line = []
+            next_line = []
+
+            for child in current:
+                line.append(child)
+                frameid_to_lineno[child._id] = line_no
+                next_line.extend(child.children)
+
+            lines.append(line)
+            line_no += 1
+            current = next_line
+
+        t2 = time.time()
+        logger.info("create lines done, took %.2f seconds", t2 - t1)
+        self.lines = lines
+        self.frameid_to_lineno = frameid_to_lineno
+
+        self.name_aggr = self.get_name_aggr(self.root_stack)
+
+    def get_name_aggr(
+        self, start_frame: Frame, names: Set[str] | None = None
+    ) -> Dict[str, List[Frame]]:
+        name = start_frame.name
+
+        result = {}
+        if names is None:
+            names = set()
+        if name not in names:
+            result[name] = [start_frame]
+
+        for child in start_frame.children:
+            name_aggr = self.get_name_aggr(child, names | set([name]))
+            for key, value in name_aggr.items():
+                result.setdefault(key, []).extend(value)
+
+        return result

flamegraph_textual/parsers/__init__.py

@@ -0,0 +1,25 @@
+import logging
+
+from flamegraph_textual.exceptions import ProfileParseException
+from flamegraph_textual.parsers.pprof_parser import ProfileParser as PprofParser
+from flamegraph_textual.parsers.stackcollapse_parser import StackCollapseParser
+
+logger = logging.getLogger(__name__)
+
+
+ALL_PARSERS = [PprofParser, StackCollapseParser]
+
+
+def choose_parser(content: bytes):
+    for p in ALL_PARSERS:
+        if p.validate(content):
+            return p
+    raise ProfileParseException("Can not match any parser")
+
+
+def parse(filecontent: bytes, filename):
+    parser_cls = choose_parser(filecontent)
+    logger.info("Using %s...", parser_cls)
+    parser = parser_cls(filename)
+    profile = parser.parse(filecontent)
+    return profile

flamegraph_textual/parsers/pprof_parser.py

@@ -0,0 +1,352 @@
+"""
+Parse golang's pprof format into flamegraph_textual models which can be rendered.
+
+Ref:
+https://github.com/google/pprof/tree/main/proto
+"""
+
+from dataclasses import dataclass, field
+import datetime
+import gzip
+import logging
+from typing import Dict, List
+
+from rich.text import Text
+
+from flamegraph_textual.models import Frame, Profile, SampleType
+
+from . import profile_pb2
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Function:
+    id: int = 0
+    filename: str = ""
+    name: str = ""
+    start_line: int = 0
+    system_name: str = ""
+
+
+@dataclass
+class Line:
+    line_no: int = 0
+    function: Function = field(default_factory=Function)
+
+
+@dataclass
+class Mapping:
+    id: int = 0
+    memory_start: int = 0
+    memory_limit: int = 0
+    file_offset: int = 0
+
+    filename: str = ""
+    build_id: str = ""
+
+    has_functions: bool | None = None
+    has_filenames: bool | None = None
+    has_line_numbers: bool | None = None
+    has_inline_frames: bool | None = None
+
+
+@dataclass
+class Location:
+    id: int = 0
+    mapping: Mapping = field(default_factory=Mapping)
+    address: int = 0
+    lines: List[Line] = field(default_factory=list)
+    is_folded: bool = False
+
+
+class PprofFrame(Frame):
+    def __init__(
+        self,
+        name,
+        _id,
+        children=None,
+        parent=None,
+        values=None,
+        root=None,
+        line=None,
+        mapping=None,
+    ) -> None:
+        super().__init__(name, _id, children, parent, values, root)
+
+        parts = self.name.split("/")
+        if len(parts) > 1:
+            self.golang_package = "/".join(parts[:-1])
+        else:
+            self.golang_package = "buildin"
+
+        self.golang_module_function = parts[-1]
+
+        self.golang_module = self.golang_module_function.split(".")[0]
+
+        self.mapping_file = ""
+        self.line = line
+        self.mapping = mapping
+
+    @property
+    def color_key(self):
+        return self.golang_module
+
+    @property
+    def display_name(self):
+        return self.golang_module_function
+
+    def render_one_frame_detail(self, frame, sample_index, sample_unit):
+        if frame._id == 0:  # root
+            total = sum([c.values[sample_index] for c in frame.children])
+            value = frame.humanize(sample_unit, total)
+            if frame.children:
+                binary_name = f"Binary: {frame.children[0].mapping.filename}"
+            else:
+                binary_name = "root"
+            detail = f"{binary_name} [b red]{value}[/b red]\n"
+            return [Text.from_markup(detail)]
+
+        value = frame.humanize(sample_unit, frame.values[sample_index])
+        line1 = f"{frame.line.function.name}: [b red]{value}[/b red]\n"
+
+        line2 = (
+            f"  [grey37]{frame.line.function.filename}, [b]line"
+            f" {frame.line.line_no}[/b][/grey37]\n"
+        )
+        return [Text.from_markup(line1), Text.from_markup(line2)]
+
+    @property
+    def title(self) -> str:
+        return self.display_name
+
+
+def unmarshal(content) -> profile_pb2.Profile:
+    if len(content) < 2:
+        raise Exception(
+            "Profile content length is too short: {} bytes".format(
+                len(content)
+            )
+        )
+    is_gzip = content[0] == 31 and content[1] == 139
+
+    if is_gzip:
+        content = gzip.decompress(content)
+
+    profile = profile_pb2.Profile()
+    profile.ParseFromString(content)
+
+    return profile
+
+
+class ProfileParser:
+    def __init__(self, filename):
+        self.filename = filename
+        # uniq id
+        self.next_id = 0
+
+        self.root = PprofFrame("root", _id=self.idgenerator())
+        # need to set PprofFrame.root for every frame
+        self.root.root = self.root
+
+        # store the pprof's string table
+        self._t = []
+        # parse cached locations, profile do not need this, so only store
+        # them on the parser
+        self.locations = []
+        self.highest = 0
+
+        self.id_store: Dict[int, Frame] = {self.root._id: self.root}
+
+    def idgenerator(self):
+        i = self.next_id
+        self.next_id += 1
+
+        return i
+
+    def s(self, index):
+        return self._t[index]
+
+    def parse_internal_data(self, pbdata):
+        self._t = pbdata.string_table
+        self.functions = self.parse_functions(pbdata.function)
+        self.mappings = self.parse_mapping(pbdata.mapping)
+        self.locations = self.parse_location(pbdata.location)
+
+    def parse(self, binary_data):
+        pbdata = unmarshal(binary_data)
+        self.parse_internal_data(pbdata)
+
+        sample_types = self.parse_sample_types(pbdata.sample_type)
+
+        root = self.root
+        root.values = [0] * len(sample_types)
+        for pbsample in pbdata.sample:
+            child_frame = self.parse_sample(pbsample)
+            if not child_frame:
+                continue
+            root.values = list(map(sum, zip(root.values, child_frame.values)))
+            root.pile_up(child_frame)
+
+        pprof_profile = Profile(
+            filename=self.filename,
+            root_stack=root,
+            highest_lines=self.highest,
+            total_sample=len(pbdata.sample),
+            sample_types=sample_types,
+            id_store=self.id_store,
+        )
+
+        # default is 0, by the doc, default should be the last one
+        if pbdata.default_sample_type:
+            pprof_profile.default_sample_type_index = (
+                pbdata.default_sample_type
+            )
+
+        pprof_profile.created_at = self.parse_created_at(pbdata.time_nanos)
+        pprof_profile.period = pbdata.period
+        pprof_profile.period_type = self.to_smaple_type(pbdata.period_type)
+
+        return pprof_profile
+
+    def parse_sample(self, sample) -> PprofFrame | None:
+        values = sample.value
+        locations = list(
+            reversed([self.locations[i] for i in sample.location_id])
+        )
+
+        my_depth = sum(len(loc.lines) for loc in locations)
+        self.highest = max(my_depth, self.highest)
+
+        current_parent = None
+        head = None
+        for location in locations:
+            for line in location.lines:
+                frame = self.line2frame(location, line, values)
+
+                if current_parent:
+                    frame.parent = current_parent
+                    current_parent.children = [frame]
+                if not head:
+                    head = frame
+
+                current_parent = frame
+
+        return head
+
+    def line2frame(self, location: Location, line: Line, values) -> PprofFrame:
+        frame = PprofFrame(
+            name=line.function.name,
+            _id=self.idgenerator(),
+            values=values,
+            root=self.root,
+            mapping=location.mapping,
+        )
+        frame.line = line
+        self.id_store[frame._id] = frame
+        return frame
+
+    def parse_location(self, pblocations):
+        parsed_locations = {}
+        for pl in pblocations:
+            loc = Location()
+            loc.id = pl.id
+            loc.mapping = self.mappings[pl.mapping_id]
+            loc.address = pl.address
+            loc.lines = self.parse_line(pl.line)
+            loc.is_folded = pl.is_folded
+            parsed_locations[loc.id] = loc
+
+        return parsed_locations
+
+    def parse_mapping(self, pbmappings):
+        mappings = {}
+        for pbm in pbmappings:
+            m = Mapping()
+            m.id = pbm.id
+            m.memory_start = pbm.memory_start
+            m.memory_limit = pbm.memory_limit
+            m.file_offset = pbm.file_offset
+            m.filename = self.s(pbm.filename)
+            m.build_id = self.s(pbm.build_id)
+            m.has_functions = pbm.has_functions
+            m.has_filenames = pbm.has_filenames
+            m.has_line_numbers = pbm.has_line_numbers
+            m.has_inline_frames = pbm.has_inline_frames
+
+            mappings[m.id] = m
+
+        return mappings
+
+    def parse_line(self, pblines) -> List[Line]:
+        lines = []
+        for pl in reversed(pblines):
+            line = Line(
+                line_no=pl.line,
+                function=self.functions[pl.function_id],
+            )
+            lines.append(line)
+        return lines
+
+    def parse_functions(self, pfs):
+        functions = {}
+        for pf in pfs:
+            functions[pf.id] = Function(
+                id=pf.id,
+                filename=self.s(pf.filename),
+                name=self.s(pf.name),
+                system_name=self.s(pf.system_name),
+                start_line=pf.start_line,
+            )
+        return functions
+
+    def parse_created_at(self, time_nanos):
+        date = datetime.datetime.fromtimestamp(
+            time_nanos / 1e9, tz=datetime.timezone.utc
+        )
+        return date
+
+    def parse_sample_types(self, sample_types):
+        result = []
+        for st in sample_types:
+            result.append(self.to_smaple_type(st))
+
+        return result
+
+    def to_smaple_type(self, st):
+        return SampleType(self.s(st.type), self.s(st.unit))
+
+    @classmethod
+    def validate(cls, content: bytes) -> bool:
+        try:
+            unmarshal(content)
+        except:  # noqa E722
+            logger.info("Error when parse content as Pprof")
+            return False
+        return True
+
+
+def get_frame_tree(root_frame):
+    """
+    only for testing and debugging
+    """
+
+    def _get_child(frame):
+        return {c.name: _get_child(c) for c in frame.children}
+
+    return {"root": _get_child(root_frame)}
+
+
+def parse_profile(binary_data, filename):
+    parser = ProfileParser(filename)
+    profile = parser.parse(binary_data)
+
+    # import ipdb; ipdb.set_trace()
+    return profile
+
+
+if __name__ == "__main__":
+    with open("tests/pprof_data/goroutine.out", "rb") as f:
+        content = f.read()
+
+    parse_profile(content, "abc")