Fix generic plot and add unit test

Author: freeman-lab

lightning/main.py

@@ -117,7 +117,7 @@ def set_host(self, host):
         self.host = host
         return self
 
-    def plot(self, type=None, **kwargs):
+    def plot(self, data=None, type=None):
         """
         Generic plotting function.
 
@@ -128,11 +128,22 @@ def plot(self, type=None, **kwargs):
 
         Most useful when providing data to custom visualizations, as opposed to the included plot types
         (e.g. lightning.scatter, lightning.line, etc.) which do automatic parsing and formatting.
+
+        Parameters
+        ----------
+        data : dict
+            Dictionary with data to plot
+
+        type : str
+            Name of plot (e.g. 'line' or 'scatter')
         """
 
         from types.plots import Generic
+
+        if not hasattr(self, 'session'):
+            self.create_session()
 
-        viz = Generic.baseplot(self.session, type=type, **kwargs)
+        viz = Generic.baseplot(self.session, type, data)
         self.session.visualizations.append(viz)
         return viz
 

lightning/types/plots.py

@@ -9,7 +9,7 @@ class Generic(Base):
 
     @staticmethod
     def clean(data):
-        return {'data': data}
+        return data
 
 
 @viztype

test/test_client.py

@@ -41,6 +41,17 @@ def test_create_scatter(self):
         assert hasattr(viz, 'id')
 
 
+    def test_create_generic(self):
+
+        series = random.randn(5,100)
+
+        viz = lightning.plot(data={"series": series}, type='line')
+        viz = lightning.plot({"series": series}, 'line')
+
+        assert isinstance(viz, Visualization)
+        assert hasattr(viz, 'id')
+
+
     def test_create_line(self):
 
         series = random.randn(5,100)