Fixes typing issues and bugs identified by mypy

.gitignore

@@ -1,16 +1,169 @@
 *.pyc
-.cache
-.pytest_cache
-.ipynb_checkpoints
 
-build
-docs/_build
 hatchet/cython_modules/libs/graphframe_modules.*.so
 hatchet/cython_modules/libs/reader_modules.*.so
 hatchet/cython_modules/*.c
 hatchet/vis/*node_modules*
 hatchet/vis/static/*_bundle*
 *package-lock.json
 
+###############################################
+# Everything from here on comes from the GitHub
+# gitignore template for Python projects
+###############################################
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
 dist/
-llnl_hatchet.egg-info/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/

hatchet/external/__init__.py

@@ -3,6 +3,8 @@
 #
 # SPDX-License-Identifier: MIT
 
+from typing import TYPE_CHECKING
+
 
 class VersionError(Exception):
     """
@@ -13,23 +15,24 @@ class VersionError(Exception):
     pass
 
 
-try:
-    import IPython
+if not TYPE_CHECKING:
+    try:
+        import IPython
 
-    # Testing IPython version
-    if int(IPython.__version__.split(".")[0]) > 7:
-        raise VersionError()
+        # Testing IPython version
+        if int(IPython.__version__.split(".")[0]) > 7:
+            raise VersionError()
 
-    from .roundtrip.roundtrip.manager import Roundtrip
+        from .roundtrip.roundtrip.manager import Roundtrip
 
-    # Refrencing Roundtrip here to resolve scope issues with import
-    Roundtrip
+        # Refrencing Roundtrip here to resolve scope issues with import
+        Roundtrip
 
-except ImportError:
-    pass
+    except ImportError:
+        pass
 
-except VersionError:
-    if IPython.get_ipython() is not None:
-        print(
-            "Warning: Roundtrip module could not be loaded. Requires jupyter notebook version <= 7.x."
-        )
+    except VersionError:
+        if IPython.get_ipython() is not None:
+            print(
+                "Warning: Roundtrip module could not be loaded. Requires jupyter notebook version <= 7.x."
+            )

hatchet/external/console.py

@@ -34,7 +34,7 @@
 import pandas as pd
 import numpy as np
 import warnings
-from typing import List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from ..util.colormaps import ColorMaps
 from ..node import Node
 
@@ -43,14 +43,19 @@ class ConsoleRenderer:
     def __init__(self, unicode: bool = False, color: bool = False) -> None:
         self.unicode = unicode
         self.color = color
-        self.visited = []
+        self.visited: List[Node] = []
+        self.colors_annotations_mapping: Optional[Union[List, Dict[str, Any]]] = None
+        self.colors: Optional[
+            Union["ConsoleRenderer.colors_enabled", "ConsoleRenderer.colors_disabled"]
+        ] = None
+        self.temporal_symbols: Dict[str, str] = {}
 
     def render(
         self,
         roots: Optional[Union[List[Node], Tuple[Node, ...]]],
         dataframe: pd.DataFrame,
         **kwargs,
-    ) -> str:
+    ) -> Union[str, bytes]:
         self.render_header = kwargs["render_header"]
 
         if self.render_header:
@@ -79,7 +84,7 @@ def render(
         self.max_value = kwargs["max_value"]
 
         if self.color:
-            self.colors = self.colors_enabled
+            self.colors = self.colors_enabled()
             # set the colormap based on user input
             self.colors.colormap = ColorMaps().get_colors(
                 self.colormap, self.invert_colormap
@@ -100,7 +105,7 @@ def render(
                 elif isinstance(self.colormap_annotations, dict):
                     self.colors_annotations_mapping = self.colormap_annotations
         else:
-            self.colors = self.colors_disabled
+            self.colors = self.colors_disabled()
 
         if isinstance(self.metric_columns, (str, tuple)):
             self.primary_metric = self.metric_columns
@@ -265,14 +270,13 @@ def render_frame(
         if node_depth < self.depth:
             # set dataframe index based on whether rank and thread are part of
             # the MultiIndex
+            df_index: Union[Tuple[Node, int, int], Tuple[Node, int], Node] = node
             if "rank" in dataframe.index.names and "thread" in dataframe.index.names:
                 df_index = (node, self.rank, self.thread)
             elif "rank" in dataframe.index.names:
                 df_index = (node, self.rank)
             elif "thread" in dataframe.index.names:
                 df_index = (node, self.thread)
-            else:
-                df_index = node
 
             node_metric = dataframe.loc[df_index, self.primary_metric]
 
@@ -326,10 +330,12 @@ def render_frame(
                 # no pattern column
                 elif self.colormap_annotations:
                     if isinstance(self.colormap_annotations, dict):
+                        assert isinstance(self.colors_annotations_mapping, Dict)
                         color_annotation = self.colors_annotations_mapping[
                             annotation_content
                         ]
                     else:
+                        assert isinstance(self.colors_annotations_mapping, List)
                         color_annotation = self.colors_annotations.colormap[
                             self.colors_annotations_mapping.index(annotation_content)
                             % len(self.colors_annotations.colormap)
@@ -441,7 +447,7 @@ def _ansi_color_for_name(self, node_name: str) -> str:
             return self.colors.bg_white_255 + self.colors.dark_gray_255
 
     class colors_enabled:
-        colormap = []
+        colormap: List[str] = []
 
         blue = "\033[34m"
         cyan = "\033[36m"
@@ -456,9 +462,7 @@ class colors_enabled:
         end = "\033[0m"
 
     class colors_disabled:
-        colormap = ["", "", "", "", "", "", ""]
+        colormap: List[str] = ["", "", "", "", "", "", ""]
 
         def __getattr__(self, key: str) -> str:
             return ""
-
-    colors_disabled = colors_disabled()

hatchet/frame.py

@@ -47,9 +47,11 @@ def __init__(self, attrs: Optional[Dict[str, Any]] = None, **kwargs) -> None:
         if "type" not in self.attrs:
             self.attrs["type"] = "None"
 
-        self._tuple_repr = None
+        self._tuple_repr: Optional[Tuple[Tuple[str, Any], ...]] = None
 
-    def __eq__(self, other: "Frame") -> bool:
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Frame):
+            return NotImplemented
         return self.tuple_repr == other.tuple_repr
 
     def __lt__(self, other: "Frame") -> bool:
@@ -69,7 +71,7 @@ def __repr__(self) -> str:
         return "Frame(%s)" % self
 
     @property
-    def tuple_repr(self) -> Tuple[Tuple[str, Any], ...]:
+    def tuple_repr(self) -> Optional[Tuple[Tuple[str, Any], ...]]:
         """Make a tuple of attributes and values based on reader."""
         if not self._tuple_repr:
             self._tuple_repr = tuple(sorted((k, v) for k, v in self.attrs.items()))