Merge 122396d into d18800a

README.md

@@ -51,12 +51,11 @@ The documentation is accessible [here](https://screamshot.readthedocs.io/en/late
 The server must be launched using --nothreading and --noreload as argument.
 ```
 # views.py in a Django project
-import asyncio
-
 from django.http import HttpResponse
 
-from screamshot import generate_bytes_img_prom
+import asyncio
 
+from screamshot import generate_bytes_img_prom
 
 def home(request):
     loop = asyncio.get_event_loop()
@@ -66,11 +65,21 @@ def home(request):
         generate_bytes_img_prom('https://www.google.fr', future))
     loop.run_until_complete(future)
 
-    print(future.result())
-    return HttpResponse('Done')
+    return HttpResponse(future.result(), content_type='image')
+``` 
+Or using the already wrapped function
+```
+# views.py in a Django project
+from django.http import HttpResponse
+
+from screamshot import generate_bytes_img__django_wrap
+
+def home(request):
+    img = generate_bytes_img__django_wrap('https://www.google.fr')
+    return HttpResponse(img, content_type='image')
 ``` 
 
 
 #### Using Gunicorn
 
-With [Gunicorn](https://gunicorn.org/) there isn't the thread related problems and we don't need to use the --nothreading and --noreload arguments.
+With [Gunicorn](https://gunicorn.org/) there isn't the thread related problems so we don't need to use the --nothreading and --noreload arguments.

screamshot/__init__.py

@@ -8,8 +8,18 @@
 __version__ = "0.1.4"
 
 
-from screamshot.generate_bytes_img_functions import generate_bytes_img, generate_bytes_img_prom
+from screamshot.generate_bytes_img_functions import (
+    generate_bytes_img,
+    generate_bytes_img_prom,
+    generate_bytes_img_django_wrap,
+)
 from screamshot.serializer import serialize, deserialize
 
 
-__all__ = ['generate_bytes_img', 'generate_bytes_img_prom', 'serialize', 'deserialize']
+__all__ = [
+    "generate_bytes_img",
+    "generate_bytes_img_prom",
+    "generate_bytes_img_django_wrap",
+    "serialize",
+    "deserialize",
+]

screamshot/generate_bytes_img_functions.py

@@ -2,7 +2,8 @@
 generate_bytes_img and generate_bytes_img_prom functions.
 """
 from typing import Any
-from asyncio.futures import Future
+from asyncio import Future
+from asyncio import get_event_loop, ensure_future
 
 from pyppeteer.page import Page
 from pyppeteer.browser import Browser
@@ -12,76 +13,76 @@
 
 def _parse_parameters(**kwargs) -> dict:
     arg_viewport = {}
-    if 'width' in kwargs:
-        arg_viewport.update({'width': kwargs.pop('width')})
-    if 'height' in kwargs:
-        arg_viewport.update({'height': kwargs.pop('height')})
+    if "width" in kwargs:
+        arg_viewport.update({"width": kwargs.pop("width")})
+    if "height" in kwargs:
+        arg_viewport.update({"height": kwargs.pop("height")})
 
-    if 'wait_until' in kwargs:
-        wait_until = kwargs.pop('wait_until')
+    if "wait_until" in kwargs:
+        wait_until = kwargs.pop("wait_until")
         if not isinstance(wait_until, list):
             wait_until = [wait_until]
     else:
-        wait_until = ['load']
+        wait_until = ["load"]
 
-    screenshot_options = {'fullPage': kwargs.get('full_page', False)}
-    if 'path' in kwargs:
-        path = kwargs.pop('path')
-        screenshot_options.update({'path': path})
+    screenshot_options = {"fullPage": kwargs.get("full_page", False)}
+    if "path" in kwargs:
+        path = kwargs.pop("path")
+        screenshot_options.update({"path": path})
 
     selector = None
-    if 'selector' in kwargs:
-        selector = kwargs.pop('selector')
+    if "selector" in kwargs:
+        selector = kwargs.pop("selector")
 
     wait_for = None
-    if 'wait_for' in kwargs:
-        wait_for = kwargs.pop('wait_for')
+    if "wait_for" in kwargs:
+        wait_for = kwargs.pop("wait_for")
 
     credentials = {}
-    if 'credentials' in kwargs:
-        credentials_data = kwargs.pop('credentials')
-        if 'username' in credentials_data and 'password' in credentials_data:
-            credentials['login'] = True
-        if 'token_in_header' in credentials_data:
-            credentials['token_in_header'] = credentials_data.pop('token_in_header')
-        credentials.update({'credentials_data': credentials_data})
+    if "credentials" in kwargs:
+        credentials_data = kwargs.pop("credentials")
+        if "username" in credentials_data and "password" in credentials_data:
+            credentials["login"] = True
+        if "token_in_header" in credentials_data:
+            credentials["token_in_header"] = credentials_data.pop("token_in_header")
+        credentials.update({"credentials_data": credentials_data})
 
     return {
-        'arg_viewport': arg_viewport,
-        'screenshot_options': screenshot_options,
-        'selector': selector,
-        'wait_for': wait_for,
-        'wait_until': wait_until,
-        'credentials': credentials,
+        "arg_viewport": arg_viewport,
+        "screenshot_options": screenshot_options,
+        "selector": selector,
+        "wait_for": wait_for,
+        "wait_until": wait_until,
+        "credentials": credentials,
     }
 
 
 async def _page_manager(browser: Browser, url: str, params: dict) -> Page:
     page = await browser.newPage()
 
-    arg_viewport = params.get('arg_viewport')
+    arg_viewport = params.get("arg_viewport")
     if arg_viewport:
         await page.setViewport(arg_viewport)
 
-    credentials = params.get('credentials')
+    credentials = params.get("credentials")
     if credentials:
-        credentials_data = credentials.get('credentials_data')
-        if credentials.get('login'):
+        credentials_data = credentials.get("credentials_data")
+        if credentials.get("login"):
             await page.authenticate(credentials_data)
-        if credentials.get('token_in_header'):
+        if credentials.get("token_in_header"):
             await page.setExtraHTTPHeaders(credentials_data)
 
-    await page.goto(url, waitUntil=params.get('wait_until'))
+    await page.goto(url, waitUntil=params.get("wait_until"))
 
-    wait_for = params.get('wait_for')
+    wait_for = params.get("wait_for")
     if wait_for:
         await page.waitForSelector(wait_for)
 
     return page
 
 
 async def _selector_manager(page: Page, params: dict) -> Any:
-    selector = params.get('selector')
+    selector = params.get("selector")
     if selector:
         return await page.querySelector(selector)
 
@@ -153,7 +154,7 @@ async def main():
 
     element = await _selector_manager(page, params)
 
-    image = await element.screenshot(options=params.get('screenshot_options'))
+    image = await element.screenshot(options=params.get("screenshot_options"))
 
     await page.close()
     return image
@@ -227,3 +228,45 @@ def home(request):
     img = await generate_bytes_img(url, **kwargs)
 
     future.set_result(img)
+
+
+def generate_bytes_img_django_wrap(url: str, **kwargs):
+    """
+    This function takes a screenshot and returns it as a `bytes` object in \
+        synchorouns mode, usable directly in django
+
+    :param url: mandatory, the website's url
+    :type url: str
+
+    :param path: optional, the path to the image output
+    :type path: str
+
+    :param width: optionnal, the window's width
+    :type width: int
+
+    :param height: optionnal, the window's height
+    :type height: int
+
+    :param selector: optionnal, CSS3 selector, item whose screenshot is taken
+    :type selector: str
+
+    :param wait_for: optionnal, CSS3 selector, item to wait before taking the screenshot
+    :type wait_for: str
+
+    :param wait_until: optionnal, define how long you wait for the page to be loaded should \
+        be either load, domcontentloaded, networkidle0 or networkidle2
+    :type wait_until: str or list(str)
+
+    :returns: the binary code of the image
+    :retype: `bytes`
+
+
+    """
+
+    loop = get_event_loop()
+    future = Future()
+
+    ensure_future(generate_bytes_img_prom(url, future, **kwargs))
+    loop.run_until_complete(future)
+
+    return future.result()