version 0.1.9 add screenshot handler

README.md

@@ -27,3 +27,12 @@ chrome.get_title("https://google.com") # Google
 + create and attach AWS API Gateway and note the url and apikey
 + `pip install chromeless`
 + write and run your method
+
+## Tips
++ if you wanna take screenshot, submit as one method like this  
+
+```python
+def get_screenshot(self,url,filename):
+    self.get(url)
+    self.save_screenshot(filename)
+```

awslambda/lambda_function.py

@@ -1,4 +1,5 @@
 from selenium.webdriver import Chrome, ChromeOptions
+from screenshot_handler import ScreenshotHandler
 import json
 import traceback
 from pprint import pprint
@@ -34,9 +35,11 @@ def lambda_handler(event, context):
         for name, func in funcs.items():
             setattr(Chrome, name, func)
         chrome = gen_chrome()
+        screenshothandler = ScreenshotHandler(chrome)
         # setattr(self, func.__name__, types.MethodType(method, self))
         method = getattr(chrome, called_name_as_method)
         result = method(*arg, **kwargs)
+        result = screenshothandler.wrap_with_screenshots(result)
         statusCode = 200
     except Exception as e:
         statusCode = 501

awslambda/screenshot_handler.py

@@ -0,0 +1,25 @@
+from selenium.webdriver import Chrome
+
+
+def save_screenshot(self, filename):
+    png = self.get_screenshot_as_png()
+    self.screenshothandler.append(filename, png)
+
+
+Chrome.save_screenshot = save_screenshot
+
+
+class ScreenshotHandler():
+    def __init__(self, chrome):
+        self.chrome = chrome
+        chrome.screenshothandler = self
+        self.screenshots = []
+
+    def append(self, filename, png):
+        self.screenshots.append((filename, png))
+
+    def wrap_with_screenshots(self, result):
+        return {
+            'screenshots': self.screenshots,
+            'origin_result': result,
+        }

awslambda/test_lambda.py

@@ -4,15 +4,16 @@
 sys.path.insert(0, '..')
 import examples
 sys.path.append('../pypi/chromeless')
-from chromeless import Chromeless, unpickle_result, dump_codes
+from chromeless import Chromeless, unpickle_result, dump_codes, exact_result_and_save_screenshots
 from selenium.webdriver import Chrome
 
 
 def request(func, *arg, **kwargs):
     examples_funcs = [
         examples.get_title,
         examples.get_list,
-        examples.get_title_letter_num
+        examples.get_title_letter_num,
+        examples.get_screenshot
     ]
     called_name_as_method = func.__name__
     stored_funcs = {func.__name__: func for func in examples_funcs}
@@ -26,7 +27,7 @@ def request(func, *arg, **kwargs):
     statusCode = d['statusCode']
     if statusCode != 200:
         raise Exception(d['body'])
-    response = unpickle_result(d['body'])
+    response = exact_result_and_save_screenshots(unpickle_result(d['body']))
     return response
 
 
@@ -50,6 +51,9 @@ def test_get(self):
         result = request(Chrome.get, "http://aws.amazon.com")
         self.assertTrue(result is None)
 
+    def test_screenshot(self):
+        result = request(Chrome.get_screenshot, "https://github.com/umihico", "screenshot.png")
+
 
 if __name__ == '__main__':
     unittest.main()

examples.py

@@ -21,6 +21,12 @@ def get_title_letter_num(self, url):
     return len(self.get_title(url))
 
 
+def get_screenshot(self, url, filename):
+    """screenshot test"""
+    self.get(url)
+    self.save_screenshot(filename)
+
+
 def run_examples():
 
     chrome = Chromeless(awsgateway_url, awsgateway_apikey)
@@ -47,6 +53,12 @@ def run_examples():
     print(result, type(result))
     # None <class 'NoneType'>
 
+    chrome = Chromeless(awsgateway_url, awsgateway_apikey)
+    chrome.attach_method(get_screenshot)
+    result = chrome.get_screenshot("https://github.com/umihico", "screenshot.png")
+    print(result, type(result))
+    # None <class 'NoneType'>
+
 
 if __name__ == '__main__':
     run_examples()

pypi/chromeless/chromeless/__init__.py

@@ -1,8 +1,11 @@
 __all__ = ["Chromeless", "dump_codes"]
 
 from .client_pickler import dump_codes, unpickle_result
+from .screenshot_client import exact_result_and_save_screenshots
 import requests
 import types
+import io
+from PIL import Image
 
 
 class Chromeless():
@@ -18,18 +21,17 @@ def __getattr__(self, methodname):
         "this is the main method"
         self.called_name_as_method = methodname
         return self._main_method_handler
-        #
-        #
-        # def method(self, *arg, **kwargs):
-        #     arg = arg or tuple()
-        #     kwargs = kwargs or dict()
-        #     return self._main_method(func, arg, kwargs)
-        # setattr(self, func.__name__, types.MethodType(method, self))
 
     def _main_method_handler(self, *arg, **kwargs):
         arg = arg or tuple()
         kwargs = kwargs or dict()
         data = dump_codes(self.called_name_as_method, arg, kwargs, self.stored_funcs)
         response = requests.post(self.gateway_url, data=data, headers=self.headers)
         body = response.text
-        return unpickle_result(body)
+        result = unpickle_result(body)
+        result = exact_result_and_save_screenshots(result)
+        return result
+
+    def save_screenshot_binary(self, filename, binarydata):
+        img = Image.open(io.BytesIO(binarydata))
+        img.save(filename)

pypi/chromeless/chromeless/screenshot_client.py

@@ -0,0 +1,13 @@
+
+
+def save_screenshots(screenshots):
+    for filename, png in screenshots:
+        with open(filename, 'wb') as f:
+            f.write(png)
+
+
+def exact_result_and_save_screenshots(result):
+    screenshots = result['screenshots']
+    save_screenshots(screenshots)
+    result = result['origin_result']
+    return result

pypi/chromeless/package_info.ppickle

@@ -4,4 +4,4 @@
  'REPONAME': None,
  'URL': None,
  'USERNAME': 'umihico',
- 'version': '0.1.7'}
+ 'version': '0.1.9'}

pypi/chromeless/version.txt

@@ -1 +1 @@
-0.1.7
+0.1.9

test_chromeless.py

@@ -2,6 +2,7 @@
 import unittest
 from apigateway_credentials import awsgateway_url, awsgateway_apikey
 from pypi.chromeless.chromeless import Chromeless
+import os
 
 
 def gen_attached_chrome(func):
@@ -35,6 +36,14 @@ def test_get(self):
         result = chrome.get("http://aws.amazon.com")
         self.assertTrue(result is None)
 
+    def test_get_screenshot(self):
+        chrome = Chromeless(awsgateway_url, awsgateway_apikey)
+        chrome.attach_method(examples.get_screenshot)
+        path = "screenshot.png"
+        result = chrome.get_screenshot("https://github.com/umihico", path)
+        self.assertTrue(result is None)
+        self.assertTrue(os.path.exists(path))
+
 
 if __name__ == '__main__':
     unittest.main()