annotate gpp/legacy/management/commands/import_old_links.py @ 333:0bf5a5677067

Import from future with_statement for Python 2.5 in import_old_links and podcasts scripts.
author Brian Neal <bgneal@gmail.com>
date Fri, 25 Feb 2011 22:03:08 +0000
parents 254db4cb6a86
children 639cfdf59167
rev   line source
bgneal@294 1 """
bgneal@294 2 import_old_links.py - For importing links from SG101 1.0 as csv files.
bgneal@294 3 """
bgneal@333 4 from __future__ import with_statement
bgneal@294 5 import csv
bgneal@294 6 import datetime
bgneal@294 7
bgneal@294 8 from django.core.management.base import LabelCommand, CommandError
bgneal@294 9 from django.contrib.auth.models import User
bgneal@294 10
bgneal@294 11 from weblinks.models import Link, Category
bgneal@294 12
bgneal@294 13
bgneal@294 14 class Command(LabelCommand):
bgneal@294 15 args = '<filename filename ...>'
bgneal@294 16 help = 'Imports weblinks from the old database in CSV format'
bgneal@294 17
bgneal@294 18 def handle_label(self, filename, **options):
bgneal@294 19 """
bgneal@294 20 Process each line in the CSV file given by filename by
bgneal@294 21 creating a new weblink object and saving it to the database.
bgneal@294 22
bgneal@294 23 """
bgneal@294 24 self.cats = {}
bgneal@294 25 try:
bgneal@294 26 self.default_user = User.objects.get(pk=2)
bgneal@294 27 except User.DoesNotExist:
bgneal@294 28 raise CommandError("Need a default user with pk=1")
bgneal@294 29
bgneal@294 30 try:
bgneal@294 31 with open(filename, "rb") as f:
bgneal@294 32 self.reader = csv.DictReader(f)
bgneal@294 33 try:
bgneal@294 34 for row in self.reader:
bgneal@294 35 self.process_row(row)
bgneal@294 36 except csv.Error, e:
bgneal@294 37 raise CommandError("CSV error: %s %s %s" % (
bgneal@294 38 filename, self.reader.line_num, e))
bgneal@294 39
bgneal@294 40 except IOError:
bgneal@294 41 raise CommandError("Could not open file: %s" % filename)
bgneal@294 42
bgneal@294 43 def get_category(self, row):
bgneal@294 44 """
bgneal@294 45 Return the Category object for the row.
bgneal@294 46
bgneal@294 47 """
bgneal@294 48 cat_id = row['cid']
bgneal@294 49 if cat_id not in self.cats:
bgneal@294 50 try:
bgneal@294 51 cat = Category.objects.get(pk=cat_id)
bgneal@294 52 except Category.DoesNotExist:
bgneal@294 53 raise CommandError("Category does not exist: %s on line %s" % (
bgneal@294 54 cat_id, self.reader.line_num))
bgneal@294 55 else:
bgneal@294 56 self.cats[cat_id] = cat
bgneal@294 57 return self.cats[cat_id]
bgneal@294 58
bgneal@294 59 def get_user(self, username):
bgneal@294 60 """
bgneal@294 61 Return the user object for the given username.
bgneal@294 62 If the user cannot be found, self.default_user is returned.
bgneal@294 63
bgneal@294 64 """
bgneal@294 65 try:
bgneal@294 66 return User.objects.get(username=username)
bgneal@294 67 except User.DoesNotExist:
bgneal@294 68 return self.default_user
bgneal@294 69
bgneal@294 70 def process_row(self, row):
bgneal@294 71 """
bgneal@294 72 Process one row from the CSV file: create an object for the row
bgneal@294 73 and save it in the database.
bgneal@294 74
bgneal@294 75 """
bgneal@294 76 link = Link(category=self.get_category(row),
bgneal@294 77 title=row['title'].decode('latin-1'),
bgneal@294 78 url=row['url'].decode('latin-1'),
bgneal@294 79 description=row['description'].decode('latin-1'),
bgneal@294 80 user=self.get_user(row['submitter']),
bgneal@294 81 date_added=datetime.datetime.strptime(row['date'], "%Y-%m-%d %H:%M:%S"),
bgneal@294 82 hits=int(row['hits']),
bgneal@294 83 is_public=True)
bgneal@294 84 link.save()