diff --git a/lightning/visualization.py b/lightning/visualization.py index 75e7b3f..ac5c316 100644 --- a/lightning/visualization.py +++ b/lightning/visualization.py @@ -1,3 +1,4 @@ +import os import requests import json import webbrowser @@ -63,11 +64,31 @@ def get_html(self): def open(self): webbrowser.open(self.session.host + '/visualizations/' + str(self.id) + '/') + def save(self, filename, width=1024, height=768, overwrite=False): + + name, ext = os.path.splitext(filename) + if len(ext) > 0: + if not ext == '.png': + raise Exception('Invalid extension %s on file name' % ext) + else: + filename += '.png' + + dir = os.path.dirname(filename) + if not os.path.isdir(dir): + raise Exception('Folder %s does not exist' % dir) + + if os.path.isfile(filename) and overwrite is False: + raise Exception('File %s exists and overwrite is set to false' % filename) + + f = open(filename, 'wb') + img = requests.get(self.get_permalink() + '/screenshot/?width=%g&height=%g' % (width, height)) + f.write(img.content) + f.close() + def delete(self): url = self.get_permalink() return requests.delete(url) - @classmethod def create(cls, session=None, data=None, images=None, type=None): url = session.host + '/sessions/' + str(session.id) + '/visualizations'