bugfix

GiulioRossetti · Aug 9, 2017 · 4981783 · 4981783
1 parent c403a2b
commit 4981783
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/ndlib/test/test_mpl_viz.py b/ndlib/test/test_mpl_viz.py
@@ -1,7 +1,6 @@
 from __future__ import absolute_import
 
 import unittest
-import os
 
 import networkx as nx
 from ndlib.viz.mpl.DiffusionPrevalence import DiffusionPrevalence
@@ -12,8 +11,8 @@
 import ndlib.models.epidemics.SIRModel as sir
 import ndlib.models.epidemics.SIModel as si
 from ndlib.viz.mpl.DiffusionTrend import DiffusionTrend
-import matplotlib
-matplotlib.use('Agg')
+
+import os
 
 
 class VizTest(unittest.TestCase):

diff --git a/ndlib/viz/mpl/ComparisonViz.py b/ndlib/viz/mpl/ComparisonViz.py
@@ -1,5 +1,10 @@
 import abc
 from bokeh.palettes import Category20_9 as cols
+import os
+import matplotlib as mpl
+if os.environ.get('DISPLAY','') == '':
+    print('no display found. Using non-interactive Agg backend')
+    mpl.use('Agg')
 import matplotlib.pyplot as plt
 import future.utils
 import past

diff --git a/ndlib/viz/mpl/DiffusionViz.py b/ndlib/viz/mpl/DiffusionViz.py
@@ -1,5 +1,10 @@
 import abc
 from bokeh.palettes import Category20_9 as cols
+import os
+import matplotlib as mpl
+if os.environ.get('DISPLAY','') == '':
+    print('no display found. Using non-interactive Agg backend')
+    mpl.use('Agg')
 import matplotlib.pyplot as plt
 import future.utils