Completed
Push — master ( 84e615...a5de1a )
by Johannes
01:08
created

Command.handle()   B

Complexity

Conditions 5

Size

Total Lines 28

Duplication

Lines 0
Ratio 0 %

Importance

Changes 4
Bugs 0 Features 1
Metric Value
c 4
b 0
f 1
dl 0
loc 28
rs 8.0894
cc 5
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
import progressbar
10
from django.apps import apps
11
from django.core.files.storage import get_storage_class
12
from django.core.management import BaseCommand, CommandError
13
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
            obj = queryset.first()
61
            do_render = True
62
            if obj:
63
                f = getattr(obj, field_name)
64
                do_render = f.field.render_variations
65
            images = queryset.values_list(field_name, flat=True).iterator()
66
            count = queryset.count()
67
68
            self.render(field, images, count, replace, do_render)
69
70
    @staticmethod
71
    def render(field, images, count, replace, do_render):
72
        pool = Pool(
73
            initializer=init_progressbar,
74
            initargs=[count]
75
        )
76
        args = [
77
            dict(
78
                file_name=file_name,
79
                do_render=do_render,
80
                variations=field.variations,
81
                replace=replace,
82
                storage=field.storage.deconstruct()[0],
83
            )
84
            for file_name in images
85
        ]
86
        pool.map(render_field_variations, args)
87
        pool.apply(finish_progressbar)
88
        pool.close()
89
        pool.join()
90
91
92
def init_progressbar(count):
93
    global BAR
94
    BAR = progressbar.ProgressBar(maxval=count, widgets=(
95
        progressbar.RotatingMarker(),
96
        ' | ', MemoryUsageWidget(),
97
        ' | CPUs: {}'.format(cpu_count()),
98
        ' | ', progressbar.AdaptiveETA(),
99
        ' | ', progressbar.Percentage(),
100
        ' ', progressbar.Bar(),
101
    ))
102
103
104
def finish_progressbar():
105
    BAR.finish()
106
107
108
def render_field_variations(kwargs):
109
    try:
110
        kwargs['storage'] = get_storage_class(kwargs['storage'])()
111
        do_render = kwargs.pop('do_render')
112
        if callable(do_render):
113
            do_render = do_render(**kwargs)
114
        if do_render:
115
            render_variations(**kwargs)
116
117
        global BAR
118
        BAR += 1
119
    except:
120
        raise Exception("".join(traceback.format_exception(*sys.exc_info())))
121