src/profiles.py

""" Helpers for loading performance profiles and extracting text from them.
"""
import json
import os
import tempfile as tf
import zipfile as zf
from typing import Literal, Optional

import hatchet as ht


class Profile:

    def __init__(self, profile_path: os.PathLike, profile_type: Literal["HPCToolkit", "CProfile", "Caliper"]):
        self.gf = self._load(profile_path, profile_type)

    def _load(self, profile_path: os.PathLike, profile_type: Literal["HPCToolkit", "CProfile", "Caliper"]) -> ht.GraphFrame:
        if profile_type == "HPCToolkit":
            toolkit_dir = profile_path[profile_path.rfind("/")+1:-4] # last dir in path, without ".zip" [:-4]
            with tf.TemporaryDirectory() as temp_dir:
                with zf.ZipFile(profile_path, 'r') as zip_ref:
                    zip_ref.extractall(temp_dir)
                    return ht.GraphFrame.from_hpctoolkit(os.path.join(temp_dir, toolkit_dir))
        elif profile_type == "CProfile":
            return ht.GraphFrame.from_cprofile(profile_path)
        elif profile_type == "Caliper":
            return ht.GraphFrame.from_caliper(profile_path)
        else:
            raise ValueError(f"Profile type {profile_type} not supported.")
        
    def profile_to_tree_str(self) -> str:
        return self.gf.tree()