Moving imports inside functions to make them lazily loaded

src/primate/plotting.py

@@ -1,14 +1,13 @@
 import numpy as np
-
-import bokeh 
-from bokeh.models import Span, Scatter, LinearAxis, Range1d, BoxAnnotation, Legend, Band, ColumnDataSource
-from bokeh.plotting import figure
-from bokeh.layouts import row
-from bokeh.models import NumeralTickFormatter
 from scipy.special import erfinv
 
-def figure_trace(info: dict, real_trace: float = None, **kwargs) -> figure:
+def figure_trace(info: dict, real_trace: float = None, **kwargs):
   """Plots the trace estimates """
+  import bokeh 
+  from bokeh.models import Span, Scatter, LinearAxis, Range1d, BoxAnnotation, Legend, Band, ColumnDataSource
+  from bokeh.plotting import figure
+  from bokeh.layouts import row
+  from bokeh.models import NumeralTickFormatter
 
   samples = np.ravel(info['convergence']['samples'])
   sample_index = np.arange(1, len(samples)+1)
@@ -75,3 +74,6 @@ def figure_trace(info: dict, real_trace: float = None, **kwargs) -> figure:
   q.add_layout(legend, "center")
   return row([p,q])
 
+def plot_trace(info: dict, real_trace: float = None, **kwargs) -> None:
+  from bokeh.plotting import show
+  show(figure_trace(info, real_trace, **kwargs))

src/primate/trace.py

@@ -237,10 +237,9 @@ def slq (
     # if verbose: te_util.print_summary(info)
     # if plot: te_plot.plot_convergence(info)
     if plot: 
-      from bokeh.plotting import show
-      from .plotting import figure_trace
-      show(figure_trace(info))
-
+      from .plotting import plot_trace
+      plot_trace(info)
+
     return (trace, info) if output_is_array else (trace[0], info)