From 097731690d766baf39b31dfe870fedb5186a8a5a Mon Sep 17 00:00:00 2001
From: Alex Tomkins <tomkins@darkzone.net>
Date: Wed, 6 Jan 2016 13:20:01 +0000
Subject: [PATCH] Allow photos to be downloaded

---
 .../commands/latest_tweets_update.py          |  9 ++++--
 .../migrations/0011_photo_image_file.py       | 19 +++++++++++
 latest_tweets/models.py                       |  1 +
 latest_tweets/utils.py                        | 32 +++++++++++++++++--
 4 files changed, 55 insertions(+), 6 deletions(-)
 create mode 100644 latest_tweets/migrations/0011_photo_image_file.py

diff --git a/latest_tweets/management/commands/latest_tweets_update.py b/latest_tweets/management/commands/latest_tweets_update.py
index 3eaaf19..2776127 100644
--- a/latest_tweets/management/commands/latest_tweets_update.py
+++ b/latest_tweets/management/commands/latest_tweets_update.py
@@ -8,7 +8,7 @@
 
 
 @transaction.atomic
-def update_user(user):
+def update_user(user, download):
     t = Twitter(auth=OAuth(
         settings.TWITTER_OAUTH_TOKEN,
         settings.TWITTER_OAUTH_SECRET,
@@ -16,7 +16,7 @@ def update_user(user):
         settings.TWITTER_CONSUMER_SECRET
     ))
     messages = t.statuses.user_timeline(screen_name=user, include_rts=True)
-    tweet_list = update_tweets(messages=messages)
+    tweet_list = update_tweets(messages=messages, download=download)
 
     # To ensure we delete any deleted tweets
     oldest_date = None
@@ -42,7 +42,10 @@ def add_arguments(self, parser):
         parser.add_argument(
             'users', metavar='user', nargs='+',
             help='Twitter username to update')
+        parser.add_argument(
+            '--download-photos', action='store_true',
+            help='Download images from photos and store locally')
 
     def handle(self, **options):
         for user in options['users']:
-            update_user(user=user)
+            update_user(user=user, download=options['download_photos'])
diff --git a/latest_tweets/migrations/0011_photo_image_file.py b/latest_tweets/migrations/0011_photo_image_file.py
new file mode 100644
index 0000000..6d6ff6c
--- /dev/null
+++ b/latest_tweets/migrations/0011_photo_image_file.py
@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('latest_tweets', '0010_photo_unique'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='photo',
+            name='image_file',
+            field=models.ImageField(blank=True, upload_to='latest_tweets/photo'),
+        ),
+    ]
diff --git a/latest_tweets/models.py b/latest_tweets/models.py
index fd4558f..5e90588 100644
--- a/latest_tweets/models.py
+++ b/latest_tweets/models.py
@@ -48,6 +48,7 @@ class Photo(models.Model):
     media_url = models.URLField('Media URL')
     large_width = models.PositiveIntegerField()
     large_height = models.PositiveIntegerField()
+    image_file = models.ImageField(upload_to='latest_tweets/photo', blank=True)
 
     class Meta:
         ordering = ('tweet', 'text_index')
diff --git a/latest_tweets/utils.py b/latest_tweets/utils.py
index aed51fb..7ad9320 100644
--- a/latest_tweets/utils.py
+++ b/latest_tweets/utils.py
@@ -1,9 +1,14 @@
 from __future__ import unicode_literals
 
 from datetime import datetime
+import hashlib
+from tempfile import TemporaryFile
 
+from django.core.files import File
 from django.utils.six.moves import html_parser
 from django.utils.timezone import utc
+from PIL import Image
+import requests
 
 from .models import Photo, Tweet
 
@@ -45,7 +50,7 @@ def tweet_html_entities(tweet, **kwargs):
     return ''.join(text)
 
 
-def tweet_photos(obj, media):
+def tweet_photos(obj, media, download):
     for photo in media:
         # Only photos
         if photo['type'] != 'photo':
@@ -63,8 +68,29 @@ def tweet_photos(obj, media):
             'large_height': int(large['h']),
         })
 
+        if download and not obj.image_file:
+            with TemporaryFile() as temp_file:
+                image_file = File(temp_file)
 
-def update_tweets(messages, tweet_entities=tweet_html_entities):
+                # Download the file
+                r = requests.get(obj.media_url, stream=True)
+                r.raise_for_status()
+
+                for chunk in r.iter_content(4096):
+                    image_file.write(chunk)
+
+                # Get Pillow to look at it
+                image_file.seek(0)
+                pil_image = Image.open(image_file)
+                image_name = '%s.%s' % (
+                    hashlib.md5(obj.media_url.encode()).hexdigest(), pil_image.format.lower())
+
+                # Save the file
+                image_file.seek(0)
+                obj.image_file.save(image_name, image_file, save=True)
+
+
+def update_tweets(messages, tweet_entities=tweet_html_entities, download=False):
     # Need to escape HTML entities
     htmlparser = html_parser.HTMLParser()
     unescape = htmlparser.unescape
@@ -114,7 +140,7 @@ def update_tweets(messages, tweet_entities=tweet_html_entities):
         })
 
         # Add any photos
-        tweet_photos(obj=obj, media=i['entities'].get('media', []))
+        tweet_photos(obj=obj, media=i['entities'].get('media', []), download=download)
 
         obj_list.append(obj)