refactor arc.get_data

john-bodley · Jan 4, 2018 · 8d01b2a · 8d01b2a
1 parent 1614b22
commit 8d01b2a
Showing 1 changed file with 5 additions and 12 deletions.
diff --git a/superset/viz.py b/superset/viz.py
@@ -2075,20 +2075,13 @@ def query_obj(self):
         return d
 
     def get_data(self, df):
-        fd = self.form_data
-        start_lat = df[fd.get('start_spatial').get('latCol')]
-        start_lon = df[fd.get('start_spatial').get('lonCol')]
-        end_lat = df[fd.get('end_spatial').get('latCol')]
-        end_lon = df[fd.get('end_spatial').get('lonCol')]
-        source_pos = list(zip(start_lon, start_lat))
-        target_pos = list(zip(end_lon, end_lat))
-
         data = []
-        for pos in list(zip(source_pos, target_pos)):
+        for item in self.process_spatial_data_obj(df):
+            pos = item.get('position')
             data.append({
-                         'sourcePosition': pos[0],
-                         'targetPosition': pos[1],
-                         })
+                'sourcePosition': pos[0],
+                'targetPosition': pos[1],
+            })
 
         return {
             'arcs': data,