importfiles.py

# coding=utf-8
from datetime import datetime
from optparse import make_option
import os
import sys

from django.contrib.auth.models import Group
from django.core.management.base import BaseCommand, CommandError
from django.core.files import File as FileWrapper

from filestorage.models import File, Directory
from organization.models import Actor


class Command(BaseCommand):
    args = '</path/to/dir>'
    help = 'Loads a directory with its children to the filstorage.'

    option_list = BaseCommand.option_list + (
        make_option('--uploader',
                    action='store',
                    dest='uploader',
                    default=None,
                    type='int',
                    help='Actor-id to be set as uploader'),
        make_option('--owner',
                    action='store',
                    dest='owner',
                    default=None,
                    type='int',
                    help='Group-id to be set as owner'),

    )

    cnt_files = 0
    cnt_directories = 0

    fs_encoding = sys.getfilesystemencoding()

    uploader = None
    owner = None

    def handle(self, *args, **options):

        print self.fs_encoding
        if len(args) != 1:
            raise CommandError('Expected a directory as an argument')
        if not os.path.exists(args[0]):
            raise CommandError('Directory %s does not exists' % args[0])

        try:
            self.uploader = Actor.objects.get(pk=options['uploader'])
            self.owner = Group.objects.get(pk=options['owner'])
        except Actor.DoesNotExist:
            raise CommandError('%s is not a valid user id' % options['uploader'])
        except Group.DoesNotExist:
            raise CommandError('%s is not a valid group id' % options['group'])

        root = Directory(name='Import %s' % datetime.now(), parent=None,
                         owner=self.owner)
        root.save()

        self._handle_directory(args[0], root)
        self.stdout.write('Successfully imported %d files in %d directories' % (self.cnt_files, self.cnt_directories))

    def _handle_directory(self, directory_path, directory):
        for root, subFolders, files in os.walk(directory_path):
            for filename in files:
                path = os.path.join(root, filename)
                with open(path, 'r') as f:
                    file_wrapper = FileWrapper(f)
                    self.cnt_files += 1
                    new_file = File(directory=directory, filename=filename,
                                    file=file_wrapper, uploader=self.uploader)
                    new_file.save()

            for folder in subFolders:
                self.cnt_directories += 1
                new_directory = Directory(name=folder, parent=directory, owner=self.owner)
                new_directory.save()
                self._handle_directory(folder, new_directory)