Completed
Push — master ( 42af00...a650c4 )
by Johannes
52s
created

stdimage.management.commands.Command.handle()   B

Complexity

Conditions 4

Size

Total Lines 23

Duplication

Lines 0
Ratio 0 %
Metric Value
dl 0
loc 23
rs 8.7972
cc 4
1
# -*- coding: utf-8 -*-
2
from __future__ import absolute_import, unicode_literals
3
4
import resource
5
import sys
6
import traceback
7
from multiprocessing import Pool, cpu_count
8
9
from django.apps import apps
10
from django.core.files.storage import get_storage_class
11
from django.core.management import BaseCommand, CommandError
12
13
import progressbar
14
from stdimage.utils import render_variations
15
16
BAR = None
17
18
19
class MemoryUsageWidget(progressbar.widgets.WidgetBase):
20
    def __call__(self, progress, data):
21
        return 'RAM: {0:10.1f} MB'.format(
22
            resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
23
        )
24
25
26
class Command(BaseCommand):
27
    help = 'Renders all variations of a StdImageField.'
28
    args = '<app.model.field app.model.field>'
29
30
    def add_arguments(self, parser):
31
        parser.add_argument('field_path',
32
                            nargs='+',
33
                            type=str,
34
                            help='<app.model.field app.model.field>')
35
        parser.add_argument('--replace',
36
                            action='store_true',
37
                            dest='replace',
38
                            default=False,
39
                            help='Replace existing files.')
40
41
    def handle(self, *args, **options):
42
        replace = options.get('replace')
43
        if len(options['field_path']):
44
            routes = options['field_path']
45
        else:
46
            routes = [options['field_path']]
47
        for route in routes:
48
            try:
49
                app_label, model_name, field_name = route.rsplit('.')
50
            except ValueError:
51
                raise CommandError("Error parsing field_path '{}'. Use format "
52
                                   "<app.model.field app.model.field>."
53
                                   .format(route))
54
            model_class = apps.get_model(app_label, model_name)
55
            field = model_class._meta.get_field(field_name)
56
57
            queryset = model_class._default_manager \
58
                .exclude(**{'%s__isnull' % field_name: True}) \
59
                .exclude(**{field_name: ''})
60
            images = queryset.values_list(field_name, flat=True).iterator()
61
            count = queryset.count()
62
63
            self.render(field, images, count, replace)
64
65
    @staticmethod
66
    def render(field, images, count, replace):
67
        pool = Pool(
68
            initializer=init_progressbar,
69
            initargs=[count]
70
        )
71
        args = [
72
            dict(
73
                file_name=file_name,
74
                variations=field.variations,
75
                replace=replace,
76
                storage=field.storage.deconstruct()[0],
77
            )
78
            for file_name in images
79
        ]
80
        pool.map(render_field_variations, args)
81
        pool.apply(finish_progressbar)
82
        pool.close()
83
        pool.join()
84
85
86
def init_progressbar(count):
87
    global BAR
88
    BAR = progressbar.ProgressBar(maxval=count, widgets=(
89
        progressbar.RotatingMarker(),
90
        ' | ', MemoryUsageWidget(),
91
        ' | CPUs: {}'.format(cpu_count()),
92
        ' | ', progressbar.AdaptiveETA(),
93
        ' | ', progressbar.Percentage(),
94
        ' ', progressbar.Bar(),
95
    ))
96
97
98
def finish_progressbar():
99
    BAR.finish()
100
101
102
def render_field_variations(kwargs):
103
    try:
104
        kwargs['storage'] = get_storage_class(kwargs['storage'])()
105
        render_variations(**kwargs)
106
        global BAR
107
        BAR += 1
108
    except:
109
        raise Exception("".join(traceback.format_exception(*sys.exc_info())))
110