From 25c62c200674209e78cec1099d6256ff08d9c5dc Mon Sep 17 00:00:00 2001
From: Ian Hellen <ianhelle@microsoft.com>
Date: Fri, 25 Jun 2021 16:49:02 -0700
Subject: [PATCH] Fixes to data_providers, data view, pivots and nbinit (#177)

* Fixes to data_providers, data view, pivots and nbinit
Replace pop with list.remove in data_providers.py
Add exception catching around Pivot instantiations in pivot_register_reader.py
Prevented push_notebook from being called before displaying control
Removed pandas option to return html schema in nbinit.py

* Updating version
---
 msticpy/_version.py                        |  2 +-
 msticpy/data/data_providers.py             |  4 ++--
 msticpy/datamodel/pivot_register_reader.py | 20 ++++++++++++++++----
 msticpy/nbtools/data_viewer.py             |  6 ++++--
 msticpy/nbtools/nbinit.py                  |  4 +---
 5 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/msticpy/_version.py b/msticpy/_version.py
index b9fd5c70e..f28fc8720 100644
--- a/msticpy/_version.py
+++ b/msticpy/_version.py
@@ -1,2 +1,2 @@
 """Version file."""
-VERSION = "1.2.1"
+VERSION = "1.2.2"
diff --git a/msticpy/data/data_providers.py b/msticpy/data/data_providers.py
index 33c7d1f8c..f0f91d7fc 100644
--- a/msticpy/data/data_providers.py
+++ b/msticpy/data/data_providers.py
@@ -330,10 +330,10 @@ def _execute_query(self, *args, **kwargs) -> Union[pd.DataFrame, Any]:
     def _check_for_time_params(self, params, missing):
         """Fall back on builtin query time if no time parameters were supplied."""
         if "start" in missing:
-            missing.pop("start", None)
+            missing.remove("start")
             params["start"] = self._query_time.start
         if "end" in missing:
-            missing.pop("end", None)
+            missing.remove("end")
             params["end"] = self._query_time.end
 
     @staticmethod
diff --git a/msticpy/datamodel/pivot_register_reader.py b/msticpy/datamodel/pivot_register_reader.py
index 117eb7154..3ae25263c 100644
--- a/msticpy/datamodel/pivot_register_reader.py
+++ b/msticpy/datamodel/pivot_register_reader.py
@@ -11,7 +11,7 @@
 import yaml
 
 from .._version import VERSION
-from ..common.exceptions import MsticpyUserConfigError
+from ..common.exceptions import MsticpyUserConfigError, MsticpyException
 from ..data.query_container import QueryContainer
 from . import entities
 from .pivot_register import PivotRegistration, create_pivot_func
@@ -20,7 +20,7 @@
 __author__ = "Ian Hellen"
 
 
-def register_pivots(
+def register_pivots(  # noqa: MC0001
     file_path: str,
     namespace: Dict[str, Any] = None,
     container: str = "other",
@@ -60,11 +60,23 @@ def register_pivots(
             )
 
         # try to import the module and retrieve the function
-        src_module = importlib.import_module(piv_reg.src_module)
+        try:
+            src_module = importlib.import_module(piv_reg.src_module)
+        except ImportError:
+            print(
+                f"Unable to add pivot functions from module '{piv_reg.src_module}'. Skipping"
+            )
+            continue
         if piv_reg.src_class:
             # if we need to get this from a class/object we need
             # to find or create one.
-            func = _get_func_from_class(src_module, namespace, piv_reg)
+            func = None
+            try:
+                func = _get_func_from_class(src_module, namespace, piv_reg)
+            except MsticpyException:
+                print(
+                    f"Unable to add pivot functions from class '{piv_reg.src_class}'. Skipping"
+                )
             if not func:
                 continue
         else:
diff --git a/msticpy/nbtools/data_viewer.py b/msticpy/nbtools/data_viewer.py
index 49b75692e..e3e6fa8d7 100644
--- a/msticpy/nbtools/data_viewer.py
+++ b/msticpy/nbtools/data_viewer.py
@@ -55,6 +55,8 @@ def __init__(
         """
         if data.empty:
             raise ValueError("No data available in 'data'")
+
+        output_notebook(hide_banner=True)
         # Drop empty columns
         data = data.dropna(axis="columns", how="all")
         self.cds = ColumnDataSource(data)
@@ -91,7 +93,6 @@ def __init__(
         self.accordion.selected_index = None
 
         self.layout = self.accordion
-        output_notebook(hide_banner=True)
 
     @property
     def filtered_data(self) -> pd.DataFrame:
@@ -134,7 +135,8 @@ def _update_data_table(self):
             print(self.data_filter.filters)
             print(len(self.filtered_data))
             print(self.filtered_data.iloc[:2])
-        push_notebook(handle=self.nb_handle)
+        if self.nb_handle:
+            push_notebook(handle=self.nb_handle)
 
     def display(self):
         """Display the widget."""
diff --git a/msticpy/nbtools/nbinit.py b/msticpy/nbtools/nbinit.py
index 1f8cf96f8..46f9e3e19 100644
--- a/msticpy/nbtools/nbinit.py
+++ b/msticpy/nbtools/nbinit.py
@@ -505,9 +505,7 @@ def _set_nb_options(namespace):
     pd.set_option("display.max_rows", 100)
     pd.set_option("display.max_columns", 50)
     pd.set_option("display.max_colwidth", 100)
-    # Set option on AML to display DataFrames with Schema
-    if os.environ.get("APPSETTING_WEBSITE_SITE_NAME") == "AMLComputeInstance":
-        pd.set_option("display.html.table_schema", True)
+
     os.environ["KQLMAGIC_LOAD_MODE"] = "silent"
     # Kqlmagic config will use AZ CLI login if available
     os.environ["KQLMAGIC_CONFIGURATION"] = "try_azcli_login=True"