Python Screenshot Generator

Recently, I needed a simple task to do: Build a website screenshot generator.

The idea is to enter any URL in an input field and show the generated screenshot in the same page.

I used Python/Django and the Selenium WebDriver to build the app. Selenium is a great tool for Web UI Automation.

Install Python Selenium with:

pip install selenium

If you want to take a screenshot using pure Python is pretty simple:

from selenium import webdriver
“”” Save a screenshot from spotify.com in current directory. “””
DRIVER = 'chromedriver'
driver = webdriver.Chrome(DRIVER)
driver.get('https://www.spotify.com')
screenshot = driver.save_screenshot('my_screenshot.png')
driver.quit()

Selenium requires a driver to interface with the browser. I’m using the Chrome Driver. Download it from here and put it in your PATH so it can be accessed from anywhere.

In my case the driver is located in /usr/local/bin/chromedriver

You can either save the screenshot on the disk:

driver.save_screenshot(IMAGE_PATH)

or get the screenshot as a binary data:

driver.get_screenshot_as_png()

We’ll support both options. So let’s start building the app:

# Create the project
django-admin.py startproject screenshot_generator
# Create the app
django-admin.py startapp app

In the settings.py file don’t forget to add the ‘app’ created, specify the templates dir and set the variables to manage the media and static data.

# screenshot_generator/settings.py
INSTALLED_APPS = [
 …
 ‘app’
]
TEMPLATES = [
    {
        ...
        'DIRS': [os.path.join(BASE_DIR, 'templates')],
        ...
    },
]
MEDIA_ROOT = os.path.join(BASE_DIR, 'media')
MEDIA_URL = '/media/'
STATIC_URL = '/static/'
STATICFILES_DIRS = (
    ('', os.path.join(BASE_DIR, 'static')),
)

In the project urls.py file, include the app urls and serve the media directory.

# screenshot_generator/urls.py
from django.urls import include, path
from django.conf import settings
from django.conf.urls.static import static
urlpatterns = [
    path('', include('app.urls')),
] + static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)

Inside our app code:

First, link the view with the url path in the urls file:

# app/urls.py
from django.urls import path
from app import views
urlpatterns = [
    path('', views.get_screenshot, name='get_screenshot'),
]

Then, create a simple HTML template with the form to enter the desired URL:

# templates/home.html
<h1>Enter URL:</h1>
<form action=”/” method=”post”>
  {% csrf_token %}
  <input type=”url” name=”url” size=”50" placeholder=”Ex: https://www.netflix.com" required>
<input type=”submit” value=”Generate screenshot”>
</form>

As you can see, the form calls a get_screenshot method that will be in our views.py . This is the method that has all the magic to generate the screenshot.

Let’s run the app and go to http://localhost:8000/ to see the preview home template:

python manage.py runserver

Finally, create the get_screenshot in our views. Let’s review it step by step:

# app/views.py
def get_screenshot(request):
    width = 1024
    height = 768

You can specify a width and height to the driver to get a custom screenshot. That’s why I have a default width and height in case these params are not specified in the URL.

Take a look at the Selenium WebDriver API documentation and see all the options.

First thing we need to do in our view is validate the request method and validate if the ‘url’ exists in that request. Also, check if the url is not null nor empty.

if request.method == 'POST' and 'url' in request.POST:
    url = request.POST.get('url', '')
    if url is not None and url != '':

Then, capture the base url, domain and url params if the user specify them:

save = False
base_url = '{0.scheme}://{0.netloc}/'.format(urlparse.urlsplit(url))
domain = urlparse.urlsplit(url)[1].split(':')[0]
params = urlparse.parse_qs(urlparse.urlparse(url).query)
if len(params) > 0:
    if ‘w’ in params: width = int(params[‘w’][0])
    if ‘h’ in params: height = int(params[‘h’][0])
# Ex: https://www.netflix.com/?w=800&h=600

After that, set the right driver, get the url and set the window size:

driver = webdriver.Chrome(DRIVER)
driver.get(url)
driver.set_window_size(width, height)

Now, check if the user has specify the save param. This variable will decide if we save the screenshot on disk or just serve it as a binary data.

if ‘save’ in params and params[‘save’][0] == ‘true’:
  # Ex: https://www.netflix.com/?save=true

If the above is true, we’ll save the screenshot in the media directory with a name that is formed joining the current timestamp and a “_image.png” string.

Pass the full path to the save_screenshot method. Also, make sure that the media directory exists:

save = True
now = str(datetime.today().timestamp())
img_dir = settings.MEDIA_ROOT
img_name = ''.join([now, '_image.png'])
full_img_path = os.path.join(img_dir, img_name)
if not os.path.exists(img_dir):
    os.makedirs(img_dir)
driver.save_screenshot(full_img_path)
screenshot = img_name

If false, just get a binary data with get_screenshot_as_png() and save it in our screenshot variable.

screenshot_img = driver.get_screenshot_as_png()
screenshot = base64.encodestring(screenshot_img)

In both cases, var_dict is the dictionary containing the variables needed for our home template. It’s all the relevant data of the screenshot generated:

var_dict = {
‘screenshot’: screenshot,
‘domain’: domain,
‘base_url’: base_url,
‘full_url’: url,
‘save’: save
}

Finally, quit the driver and render the home template:

driver.quit()
return render(request, ‘home.html’, var_dict)

# Final app/views.py:
import base64
import os
import urllib.parse as urlparse
from django.shortcuts import render
from django.http import HttpResponse
from django.conf import settings
from datetime import datetime
from selenium import webdriver
DRIVER = 'chromedriver'
def get_screenshot(request):
  """
  Take a screenshot and return a png file based on the url.
  """
  width = 1024
  height = 768
if request.method == 'POST' and 'url' in request.POST:
    url = request.POST.get('url', '')
      if url is not None and url != '':
        save = False
        base_url = '{0.scheme}://{0.netloc}/'.format(urlparse.urlsplit(url))
        domain = urlparse.urlsplit(url)[1].split(':')[0]
        params = urlparse.parse_qs(urlparse.urlparse(url).query)
        if len(params) > 0:
          if 'w' in params: width = int(params['w'][0])
          if 'h' in params: height = int(params['h'][0])
        driver = webdriver.Chrome(DRIVER)
        driver.get(url)
        driver.set_window_size(width, height)
if 'save' in params and params['save'][0] == 'true':
          save = True
          now = str(datetime.today().timestamp())
          img_dir = settings.MEDIA_ROOT
          img_name = ''.join([now, '_image.png'])
          full_img_path = os.path.join(img_dir, img_name)
          if not os.path.exists(img_dir):
            os.makedirs(img_dir)
          driver.save_screenshot(full_img_path)
          screenshot = img_name
        else:
          screenshot_img = driver.get_screenshot_as_png()
          screenshot = base64.encodestring(screenshot_img)
var_dict = {
          'screenshot': screenshot,
          'domain': domain,
          'base_url': base_url,
          'full_url': url,
          'save': save
        }
driver.quit()    
        return render(request, 'home.html', var_dict)
    else:
      return render(request, 'home.html')

Wait a second, we didn’t explain var_dict in detail and what we are passing to our home template.

When saving the screenshot on disk, just pass the image name in the ‘screenshot’ template tag and a template tag called save.

When getting a binary data (In this case an image), it’s necessary to encode/decode the image, so we can pass it in our dictionary as a string. We can’t pass a binary data to the render method.

In Python there is a module called base64 that help us to achieve this easily:

import base64
# Encode the image.
screenshot = driver.get_screenshot_as_png()
image_64_encode = base64.encodestring(screenshot)

For decoding, create a file app_extras.py in the dir app/templatetags/ and register the custom template tag decode_image.

# app/templatetags/app_extras.py
import base64
from django import template
register = template.Library()
@register.filter()
def decode_image(encoded_image):
    return "data:image/png;base64,%s" % encoded_image.decode("utf8")

Let’s add some CSS love to improve the look and feel of our page. Create a styles.css in the static/css folder:

# static/css/styles.css

body {
background-color: #ECF0F1;
color: #333333;
font-family: Arial, Helvetica, sans-serif;
margin: 0 auto;
padding: 6px;
}

a {
color: #246099;
}

div {
text-align: center
}

h1 {
font-size: 48px;
margin-bottom: 12px;
}

input, button {
margin: 12px 0;
border: 2px solid #246099;
}

input {
color: #4D4D4D;
padding: 14px;
}

button {
color: #FFF;
cursor: pointer;
background-color: #246099;
padding: 12px;
font-size: 16px;
}

small {
font-size: small;
}

Now, update the home.html template to show the screenshot and its relevant data:

# Final templates/home.html
{% load app_extras %}
{% load static %}
<!DOCTYPE html>
<html lang="en">
<head>
    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
    <meta charset="utf-8">
    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <link rel="stylesheet" type="text/css" href="{% static "css/styles.css" %}" />
    <title>Python Screenshot Generator</title>
</head>
<body>
<div>
  <h1>Python Screenshot Generator</h1>
  <form action="/" method="post">
    {% csrf_token %}
    <input type="url" name="url" size="50" value="{{ full_url }}" placeholder="Enter URL: https://www.netflix.com" required>
    <button type="submit">Generate screenshot</button>
  </form>
  <small>Optional params <strong>w, h</strong> for screenshot size and <strong>save</strong> for saving the screenshot in your media dir. <br /> Ex: https://nodejs.org/?w=800&h=600&save=true</small>
  
  {% if screenshot %}
  <h2>Screenshot for <a href="{{ base_url }}" target="_blank">{{ domain }}</a></h2>
    {% if save %}
      <img src="{{ MEDIA_URL }}{{ screenshot }}">
    {% else %}
      <img src="{{ screenshot|decode_image }}">
    {% endif %}
  {% endif %}
</div>
</body>
</html>

Let’s explain the above code:

If the screenshot tag exists, show the HTML tags. There is no need to show them if a screenshot has not been generated.

If we save the screenshot, then show it passing the media URL plus the image name to the src. It will end in something like:

<img src=”/media/1484738370.392261_image.png”>

Otherwise, we show the image passing the encoded screenshot and decode it with the template filter we created.

It will end in something like:
<img src=”data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABAAAAAJPCAIAAABtoTFXAAAgAElEQVR4nOS9W6xtzXIe9FV1jzHn&#10…>

Let’s run our final app to see it in action:

python manage.py runserver

As an example, I’m using the Node.js website url. I specify a width, height and save the screenshot on disk:

https://nodejs.org/?w=800&h=600&save=true

You can find all the source code of this python-screenshot-generator on GitHub.

I hope this Python/Django example can help someone to understand how to take a screenshot and show it in a webpage.

Enjoy

Ronny