Fix error when latest version 3.0.3 of bokeh is installed (#630)

* Fix error when latest version 3.0.3 of bokeh is installed * fix tests import * fix missing properties * reformatting graphs for plotting * Fix graph nodes attributes * Fix numbers and captions for nodes * Add bokeh 2.4.3 compatibility * Added documentation to figure_dimension * Fixed tests and use in-package imports * More formatting fixes * isort run * Formatting pydocstyle * formatting fixes * Fixing lint issues * Fixing mypy errors for Bokeh 3.0 Adding mypy to pre-commit Suppressing mypy annotation-unchecked in github action * Re-wrote parameter wrapper as decorator - tested with Bokeh v2 and v3 * Reverting accidentally modified conf.py * Updating bokeh reqs, bumping version to 2.3.2 --------- Co-authored-by: Claudiu Toma <ctoma@microsoft.com> Co-authored-by: Ian Hellen <ianhelle@microsoft.com>
microsoft · Mar 14, 2023 · 8d83f69 · 8d83f69
1 parent 193cc74
commit 8d83f69
Show file tree

Hide file tree

Showing 22 changed files with 340 additions and 428 deletions.
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
@@ -156,7 +156,7 @@ jobs:
             ${{ runner.os }}-build-mypy
       - name: mypy
         run: |
-          mypy --ignore-missing-imports --follow-imports=silent --show-column-numbers --junit-xml junit/mypy-test-${{ matrix.python-version }}-results.xml msticpy
+          mypy --ignore-missing-imports --follow-imports=silent --show-column-numbers --show-error-end --show-error-context --disable-error-code annotation-unchecked --junit-xml junit/mypy-test-${{ matrix.python-version }}-results.xml msticpy
         if: ${{ always() }}
       - name: Upload mypy test results
         uses: actions/upload-artifact@v2

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -41,6 +41,10 @@ repos:
       - id: pydocstyle
         args:
           - --convention=numpy
+  # - repo: https://github.com/pre-commit/mirrors-mypy
+  #   rev: v1.0.1
+  #   hooks:
+  #     - id: mypy
   - repo: local
     hooks:
       - id: check_reqs_all

diff --git a/.pylintrc b/.pylintrc
@@ -483,4 +483,4 @@ known-third-party=enchant
 
 # Exceptions that will emit a warning when being caught. Defaults to
 # "Exception".
-overgeneral-exceptions=Exception
+overgeneral-exceptions=builtins.Exception
diff --git a/conda/conda-reqs.txt b/conda/conda-reqs.txt
@@ -10,7 +10,7 @@ azure-mgmt-network>=2.7.0
 azure-mgmt-resource>=16.1.0
 azure-storage-blob>=12.5.0
 beautifulsoup4>=4.0.0
-bokeh>=1.4.0, <3.0.0
+bokeh>=1.4.0, <=3.0.2
 cryptography>=3.1
 deprecated>=1.2.4
 dnspython>=2.0.0, <3.0.0

diff --git a/docs/notebooks/ContiLeaksAnalysis.ipynb b/docs/notebooks/ContiLeaksAnalysis.ipynb
diff --git a/msticpy/_version.py b/msticpy/_version.py
@@ -1,2 +1,2 @@
 """Version file."""
-VERSION = "2.3.1"
+VERSION = "2.3.2"
diff --git a/msticpy/datamodel/entities/process.py b/msticpy/datamodel/entities/process.py
@@ -152,7 +152,11 @@ def description_str(self) -> str:
     def name_str(self) -> str:
         """Return Entity Name."""
         if self.ImageFile:
-            return f"{self.ImageFile.name_str}[pid:{self.ProcessId}]"
+            if isinstance(self.ImageFile, str):
+                return f"{self.ImageFile}[pid:{self.ProcessId}]"
+            if isinstance(self.ImageFile, File):
+                return f"{self.ImageFile.name_str}[pid:{self.ProcessId}]"
+            return super().name_str
         return self.ImageFile.name_str if self.ImageFile else super().name_str
 
     _entity_schema = {

diff --git a/msticpy/vis/data_viewer.py b/msticpy/vis/data_viewer.py
@@ -10,7 +10,7 @@
 import ipywidgets as widgets
 import pandas as pd
 from bokeh.io import output_notebook, push_notebook, show
-from bokeh.models import (
+from bokeh.models import (  # type: ignore[attr-defined]
     BooleanFilter,
     CDSView,
     ColumnDataSource,

diff --git a/msticpy/vis/entity_graph_tools.py b/msticpy/vis/entity_graph_tools.py
@@ -12,7 +12,7 @@
 import pandas as pd
 from bokeh.io import output_notebook, show
 from bokeh.layouts import column
-from bokeh.models import Circle, HoverTool, Label, LayoutDOM
+from bokeh.models import Circle, HoverTool, Label, LayoutDOM  # type: ignore
 from bokeh.plotting import figure, from_networkx
 
 from .._version import VERSION
@@ -24,13 +24,17 @@
 from ..nbtools.security_alert import SecurityAlert
 from ..vis.timeline import display_timeline
 from ..vis.timeline_duration import display_timeline_duration
+from .figure_dimension import bokeh_figure
 
 __version__ = VERSION
 __author__ = "Pete Bryan"
 
 req_alert_cols = ["DisplayName", "Severity", "AlertType"]
 req_inc_cols = ["id", "name", "properties.severity"]
 
+# wrap figure function to handle v2/v3 parameter renaming
+figure = bokeh_figure(figure)  # type: ignore[assignment, misc]
+
 
 @export
 class EntityGraph:
@@ -347,7 +351,7 @@ def to_df(self) -> pd.DataFrame:
         return tl_df
 
     def _add_incident_or_alert_node(self, incident: Union[Incident, Alert, None]):
-        """Check what type of entity is passed in and creates relevent graph."""
+        """Check what type of entity is passed in and creates relevant graph."""
         if isinstance(incident, Incident):
             self._add_incident_node(incident)
         elif isinstance(incident, Alert):
@@ -485,24 +489,43 @@ def plot_entitygraph(  # pylint: disable=too-many-locals
         )
     )
 
+    entity_graph_for_plotting = nx.Graph()
+    index_node = 0
+    rev_index = {}
+    fwd_index = {}
+    node_attributes = {}
+    for node_key in entity_graph.nodes:
+        entity_graph_for_plotting.add_node(index_node)
+        rev_index[node_key] = index_node
+        fwd_index[index_node] = node_key
+        node_attributes[index_node] = entity_graph.nodes[node_key]
+        index_node += 1
+
+    nx.set_node_attributes(entity_graph_for_plotting, node_attributes)
+
+    for source_node, target_node in entity_graph.edges:
+        entity_graph_for_plotting.add_edge(
+            rev_index[source_node], rev_index[target_node]
+        )
+
     graph_renderer = from_networkx(
-        entity_graph, nx.spring_layout, scale=scale, center=(0, 0)
+        entity_graph_for_plotting, nx.spring_layout, scale=scale, center=(0, 0)
     )
 
     graph_renderer.node_renderer.glyph = Circle(
         size=node_size, fill_color="node_color", fill_alpha=0.5
     )
     # pylint: disable=no-member
-    plot.renderers.append(graph_renderer)
+    plot.renderers.append(graph_renderer)  # type: ignore[attr-defined]
 
     # Create labels
-    for name, pos in graph_renderer.layout_provider.graph_layout.items():
+    for index, pos in graph_renderer.layout_provider.graph_layout.items():
         label = Label(
             x=pos[0],
             y=pos[1],
             x_offset=5,
             y_offset=5,
-            text=name,
+            text=fwd_index[int(index)],
             text_font_size=font_pnt,
         )
         plot.add_layout(label)

diff --git a/msticpy/vis/figure_dimension.py b/msticpy/vis/figure_dimension.py
@@ -0,0 +1,83 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+"""figure_dimension - helps set the width and height properties of a figure for plotting."""
+from functools import wraps
+from typing import Any, Callable
+
+from bokeh.plotting import figure
+
+from .._version import VERSION
+from ..common.utility import export
+
+__version__ = VERSION
+__author__ = "Claudiu Toma"
+
+
+@export
+def set_figure_size(fig: figure, width: int, height: int) -> figure:
+    """Set the figure size.
+
+    Args:
+        fig (figure): initial figure
+        width (int): width dimension
+        height (int): height dimension
+
+    Returns
+    -------
+        figure: figure with correct width and height
+
+    """
+    if hasattr(figure(), "height"):
+        setattr(fig, "height", height)
+    if hasattr(figure(), "plot_height"):
+        setattr(fig, "plot_height", height)
+    if hasattr(figure(), "width"):
+        setattr(fig, "width", width)
+    if hasattr(figure(), "plot_width"):
+        setattr(fig, "plot_width", width)
+    return fig
+
+
+_BOKEH_3_FIG_PARAMS = {
+    "plot_height": "height",
+    "plot_width": "width",
+}
+_BOKEH_2_FIG_PARAMS = {val: key for key, val in _BOKEH_3_FIG_PARAMS.items()}
+
+
+def bokeh_figure(func) -> Callable[[Any], Any]:
+    """
+    Wrap bokeh 'figure' function for version-agnostic param naming.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to wrap
+
+    Returns
+    -------
+    Callable :
+        The wrapped figure function.
+
+    """
+
+    @wraps(func)
+    def set_figure_size_params(*args, **kwargs):
+        """Re-write parameters names for Bokeh version."""
+        if func == figure:
+            param_mapper = (
+                _BOKEH_3_FIG_PARAMS
+                if hasattr(func(), "height")
+                else _BOKEH_2_FIG_PARAMS
+            )
+
+            func_kwargs = {
+                param_mapper.get(param, param): value for param, value in kwargs.items()
+            }
+            return func(*args, **func_kwargs)
+        return func(*args, **kwargs)
+
+    return set_figure_size_params
diff --git a/msticpy/vis/matrix_plot.py b/msticpy/vis/matrix_plot.py
@@ -16,10 +16,14 @@
 
 from .._version import VERSION
 from ..common.utility import check_kwargs
+from .figure_dimension import bokeh_figure
 
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
+# wrap figure function to handle v2/v3 parameter renaming
+figure = bokeh_figure(figure)  # type: ignore[assignment, misc]
+
 
 @attr.s(auto_attribs=True)
 class PlotParams:
@@ -167,12 +171,12 @@ def plot_matrix(data: pd.DataFrame, **kwargs) -> LayoutDOM:
 
     plot = figure(
         title=param.title,
-        plot_width=param.width,
-        plot_height=param.height,
         x_range=x_range,
         y_range=y_range,
         tools=["wheel_zoom", "box_zoom", "pan", "reset", "save"],
         toolbar_location="above",
+        width=param.width,
+        height=param.height,
     )
 
     tool_tips = [

diff --git a/msticpy/vis/network_plot.py b/msticpy/vis/network_plot.py
@@ -8,7 +8,7 @@
 
 import networkx as nx
 from bokeh.io import output_notebook
-from bokeh.models import (
+from bokeh.models import (  # type: ignore[attr-defined]
     BoxSelectTool,
     Circle,
     EdgesAndLinkedNodes,
@@ -25,10 +25,14 @@
 from typing_extensions import Literal
 
 from .._version import VERSION
+from .figure_dimension import bokeh_figure
 
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
+# wrap figure function to handle v2/v3 parameter renaming
+figure = bokeh_figure(figure)  # type: ignore[assignment, misc]
+
 
 GraphLayout = Union[
     Callable[[Any], Dict[str, Tuple[float, float]]],
@@ -142,14 +146,35 @@ def plot_nx_graph(
         height=height,
     )
 
-    graph_layout = _get_graph_layout(nx_graph, layout, **kwargs)
-    graph_renderer = from_networkx(nx_graph, graph_layout, scale=scale, center=(0, 0))
+    nx_graph_for_plotting = nx.Graph()
+    index_node = 0
+    rev_index = {}
+    fwd_index = {}
+    node_attributes = {}
+    for node_key in nx_graph.nodes:
+        nx_graph_for_plotting.add_node(index_node)
+        rev_index[node_key] = index_node
+        fwd_index[index_node] = node_key
+        node_attributes[index_node] = nx_graph.nodes[node_key]
+        index_node += 1
+
+    nx.set_node_attributes(nx_graph_for_plotting, node_attributes)
+
+    for source_node, target_node in nx_graph.edges:
+        nx_graph_for_plotting.add_edge(rev_index[source_node], rev_index[target_node])
+
+    graph_layout = _get_graph_layout(nx_graph_for_plotting, layout, **kwargs)
+
+    graph_renderer = from_networkx(
+        nx_graph_for_plotting, graph_layout, scale=scale, center=(0, 0)
+    )
     _create_edge_renderer(graph_renderer, edge_color=edge_color)
     _create_node_renderer(graph_renderer, node_size, "node_color")
 
     graph_renderer.selection_policy = NodesAndLinkedEdges()
     graph_renderer.inspection_policy = EdgesAndLinkedNodes()
-    plot.renderers.append(graph_renderer)  # pylint: disable=no-member
+    # pylint: disable=no-member
+    plot.renderers.append(graph_renderer)  # type: ignore[attr-defined]
 
     hover_tools = [
         _create_node_hover(source_attrs, target_attrs, [graph_renderer.node_renderer])
@@ -162,13 +187,13 @@ def plot_nx_graph(
 
     # Create labels
     # pylint: disable=no-member
-    for name, pos in graph_renderer.layout_provider.graph_layout.items():
+    for index, pos in graph_renderer.layout_provider.graph_layout.items():
         label = Label(
             x=pos[0],
             y=pos[1],
             x_offset=5,
             y_offset=5,
-            text=name,
+            text=fwd_index[int(index)],
             text_font_size=font_pnt,
         )
         plot.add_layout(label)
@@ -307,7 +332,7 @@ def plot_entity_graph(
         size=node_size, fill_color="node_color", fill_alpha=0.5
     )
     # pylint: disable=no-member
-    plot.renderers.append(graph_renderer)
+    plot.renderers.append(graph_renderer)  # type: ignore[attr-defined]
 
     # Create labels
     for name, pos in graph_renderer.layout_provider.graph_layout.items():
@@ -316,7 +341,7 @@ def plot_entity_graph(
             y=pos[1],
             x_offset=5,
             y_offset=5,
-            text=name,
+            text=str(name),
             text_font_size=font_pnt,
         )
         plot.add_layout(label)