Mercurial > public > sg101

annotate legacy/management/commands/import_old_users.py @ 588:c31841b52d62

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Add docutils to requirements.txt.
author Brian Neal <bgneal@gmail.com>
date Thu, 10 May 2012 20:30:46 -0500
parents ee87ea74d46b
children 9e803323a0d0
rev   line source
bgneal@290 1 """
bgneal@290 2 import_old_users.py - For importing users from SG101 1.0 as csv files.
bgneal@290 3 """
bgneal@290 4 from __future__ import with_statement
bgneal@290 5 import csv
bgneal@290 6 import optparse
bgneal@290 7 import re
bgneal@290 8 import sys
bgneal@290 9 from datetime import datetime
bgneal@290 10
bgneal@290 11 import postmarkup
bgneal@290 12
bgneal@290 13 from django.core.management.base import LabelCommand, CommandError
bgneal@290 14 from django.contrib.auth.models import User
bgneal@290 15
bgneal@290 16 import bio.models
bgneal@290 17 from legacy.phpbb import unphpbb
bgneal@290 18 from legacy.html2md import MarkdownWriter
bgneal@290 19
bgneal@290 20 TIME_ZONES = {
bgneal@290 21 '-5': 'US/Eastern',
bgneal@290 22 '-6': 'US/Central',
bgneal@290 23 '-7': 'US/Mountain',
bgneal@290 24 '-8': 'US/Pacific',
bgneal@290 25 }
bgneal@290 26 USERNAME_RE = re.compile(r'^[\w.@+-]+$')
bgneal@290 27 USERNAME_LEN = (1, 30) # min & max length values
bgneal@290 28
bgneal@290 29
bgneal@290 30 def _valid_username(username):
bgneal@290 31 """
bgneal@290 32 Return true if the username is valid.
bgneal@290 33 """
bgneal@290 34 return (USERNAME_LEN[0] <= len(username) <= USERNAME_LEN[1] and
bgneal@290 35 USERNAME_RE.match(username))
bgneal@290 36
bgneal@290 37
bgneal@290 38 def _break_name(name):
bgneal@290 39 """
bgneal@290 40 Break name into a first and last name.
bgneal@290 41 Return a 2-tuple of first_name, last_name.
bgneal@290 42 """
bgneal@290 43 parts = name.split()
bgneal@290 44 n = len(parts)
bgneal@290 45 if n == 0:
bgneal@290 46 t = '', ''
bgneal@290 47 elif n == 1:
bgneal@290 48 t = parts[0], ''
bgneal@290 49 else:
bgneal@290 50 t = ' '.join(parts[:-1]), parts[-1]
bgneal@290 51 return t[0][:USERNAME_LEN[1]], t[1][:USERNAME_LEN[1]]
bgneal@290 52
bgneal@290 53
bgneal@290 54 class Command(LabelCommand):
bgneal@290 55 args = '<filename filename ...>'
bgneal@290 56 help = 'Imports users from the old database in CSV format'
bgneal@290 57 option_list = LabelCommand.option_list + (
bgneal@290 58 optparse.make_option("-s", "--super-user",
bgneal@290 59 help="Make the user with this name a superuser"),
bgneal@290 60 optparse.make_option("-a", "--anon-user",
bgneal@290 61 help="Make the user with this name the anonymous user "
bgneal@290 62 "[default: Anonymous]"),
bgneal@290 63 optparse.make_option("-p", "--progress", action="store_true",
bgneal@290 64 help="Output a . after every 20 users to show progress"),
bgneal@290 65 )
bgneal@290 66 bb_parser = postmarkup.create(use_pygments=False, annotate_links=False)
bgneal@290 67 md_writer = MarkdownWriter()
bgneal@290 68
bgneal@290 69 def handle_label(self, filename, **options):
bgneal@290 70 """
bgneal@290 71 Process each line in the CSV file given by filename by
bgneal@290 72 creating a new user and profile.
bgneal@290 73
bgneal@290 74 """
bgneal@290 75 self.superuser = options.get('super_user')
bgneal@290 76 self.anonymous = options.get('anon_user')
bgneal@290 77 if self.anonymous is None:
bgneal@290 78 self.anonymous = 'Anonymous'
bgneal@290 79 self.show_progress = options.get('progress')
bgneal@290 80
bgneal@290 81 if self.superuser == self.anonymous:
bgneal@290 82 raise CommandError("super-user name should not match anon-user")
bgneal@290 83
bgneal@290 84 try:
bgneal@290 85 with open(filename, "rb") as f:
bgneal@290 86 self.reader = csv.DictReader(f)
bgneal@290 87 num_rows = 0
bgneal@290 88 try:
bgneal@290 89 for row in self.reader:
bgneal@290 90 self.process_row(row)
bgneal@290 91 num_rows += 1
bgneal@291 92 if self.show_progress and num_rows % 20 == 0:
bgneal@290 93 sys.stdout.write('.')
bgneal@290 94 sys.stdout.flush()
bgneal@290 95 except csv.Error, e:
bgneal@290 96 raise CommandError("CSV error: %s %s %s" % (
bgneal@290 97 filename, self.reader.line_num, e))
bgneal@290 98
bgneal@290 99 print
bgneal@290 100
bgneal@290 101 except IOError:
bgneal@290 102 raise CommandError("Could not open file: %s" % filename)
bgneal@290 103
bgneal@290 104 def process_row(self, row):
bgneal@290 105 """
bgneal@290 106 Process one row from the CSV file: create a user and user profile for
bgneal@290 107 the row and save it in the database.
bgneal@290 108
bgneal@290 109 """
bgneal@290 110 row = dict((k, v if v != 'NULL' else '') for k, v in row.iteritems())
bgneal@290 111
bgneal@290 112 if not _valid_username(row['username']):
bgneal@290 113 print "Skipping import of %s; invalid username" % row['username']
bgneal@290 114 return
bgneal@290 115
bgneal@290 116 n = User.objects.filter(username=row['username']).count()
bgneal@290 117 if n > 0:
bgneal@290 118 print "Skipping import of %s; user already exists" % row['username']
bgneal@290 119 return
bgneal@290 120
bgneal@290 121 first_name, last_name = _break_name(row['name'])
bgneal@290 122 is_superuser = self.superuser == row['username']
bgneal@290 123 is_anonymous = self.anonymous == row['username']
bgneal@290 124
bgneal@290 125 u = User(id=int(row['user_id']),
bgneal@290 126 username=row['username'],
bgneal@290 127 first_name=first_name,
bgneal@290 128 last_name=last_name,
bgneal@290 129 email=row['user_email'],
bgneal@290 130 password=row['user_password'] if row['user_password'] else None,
bgneal@290 131 is_staff=is_superuser,
bgneal@290 132 is_active=True if not is_anonymous else False,
bgneal@290 133 is_superuser=is_superuser,
bgneal@290 134 last_login=datetime.fromtimestamp(int(row['user_lastvisit'])),
bgneal@290 135 date_joined=datetime.strptime(row['user_regdate'], "%b %d, %Y"))
bgneal@290 136
bgneal@290 137 if is_anonymous:
bgneal@290 138 u.set_unusable_password()
bgneal@290 139
bgneal@290 140 u.save()
bgneal@290 141
bgneal@290 142 p = u.get_profile()
bgneal@294 143 p.location = row['user_from'].decode('latin-1')
bgneal@294 144 p.occupation = row['user_occ'].decode('latin-1')
bgneal@294 145 p.interests = row['user_interests'].decode('latin-1')
bgneal@290 146 p.profile_text = u''
bgneal@290 147 p.hide_email = True if row['user_viewemail'] != '1' else False
bgneal@290 148 p.signature = self.to_markdown(row['user_sig']) if row['user_sig'] else u''
bgneal@290 149 p.time_zone = TIME_ZONES.get(row['user_timezone'], 'US/Pacific')
bgneal@290 150 p.use_24_time = False
bgneal@290 151 p.forum_post_count = int(row['user_posts'])
bgneal@290 152 p.status = bio.models.STA_ACTIVE if p.forum_post_count > 10 else bio.models.STA_STRANGER
bgneal@290 153 p.status_date = datetime.now()
bgneal@290 154 p.update_date = p.status_date
bgneal@290 155 p.save()
bgneal@290 156
bgneal@290 157 def to_html(self, s):
bgneal@290 158 return self.bb_parser.render_to_html(unphpbb(s), cosmetic_replace=False)
bgneal@290 159
bgneal@290 160 def to_markdown(self, s):
bgneal@290 161 self.md_writer.reset()
bgneal@290 162 self.md_writer.feed(self.to_html(s))
bgneal@290 163 return self.md_writer.markdown()