Hello Developers. Today, we will be programmatically saving screenshots of websites (generated image) in the Django model. We will use Python’s selenium to take screenshots and will save in in a Django database within the model field.

Install prerequisites

Install the below stuff in your Python virtual environment and Django project.

pip install django selenium chromedriver-binary-auto Pillow

chromedriver-binary-auto is to provide an executable webdriver binary to selenium.

Config models.py

Our model would look something like below

from django.db import models

# Our model is here.
class Website(models.Model):
    name = models.CharField(max_length=100)
    url = models.URLField()
    description = models.TextField(blank=True)
    screenshot = models.ImageField(upload_to='screenshots', blank=True, null=True)

    def __str__(self):
        return self.name

    def save(self, *args, **kwargs):
        if not self.screenshot:
            import chromedriver_binary
            from selenium import webdriver

            options = webdriver.ChromeOptions()
            options.add_argument('headless')
            options.add_argument('window-size=1200x600')
            driver = webdriver.Chrome(options=options)
            driver.get(self.url)
            image_path = 'screenshots/{}.png'.format(self.name)
            driver.save_screenshot(f"media/{image_path}")
            self.screenshot =  image_path
            driver.quit()

        return super(Website, self).save(*args, **kwargs)

We have imported chromedriver_binary and selenium in the save method. Whenever the model save is triggered, the selenium opens the website in the headless mode, saves it in the path of media/screenshots/<website_name>.

And after saving it, it returns the image_path to self.screenshot so that the path gets linked. And gets saved.

Add media in settings.py

In your settings.py file, you have to add MEDIA_URL and MEDIA_ROOT, for now we will use the popular name for media, basically /media/

MEDIA_URL = '/media/'
MEDIA_ROOT = BASE_DIR / 'media'

And create a directory named media. And another directory inside media, screenshots.

media/screenshots

Add media path to urls.py

Now, you have to add MEDIA_ROOT and MEDIA_URL in your static path.

from django.contrib import admin
from django.urls import path

from django.conf.urls.static import static
from django.conf import settings

urlpatterns = [
    path('admin/', admin.site.urls),
]

# Adding MEDIA PATH
urlpatterns += static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)

Here’s a well-functioning project’s repo related to the task we are doing that you’d love to check out: github/save-generated-image-django

Good to go

Now, you can register the Website model in admin.py, and save your model in Django admin to test it. And you’ll get the screenshot of the website you saved.