web/lib/django_extensions/management/commands/runscript.py
author ymh <ymh.work@gmail.com>
Fri, 22 Jan 2010 18:23:34 +0100
changeset 11 f236caaceb43
parent 3 526ebd3988b0
permissions -rw-r--r--
add pois
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
3
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     1
from django.core.management.base import BaseCommand
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     2
from django.core.management.color import no_style
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     3
from optparse import make_option
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     4
import sys
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     5
import os
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     6
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     7
try:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     8
    set
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     9
except NameError:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    10
    from sets import Set as set   # Python 2.3 fallback
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    11
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    12
class Command(BaseCommand):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    13
    option_list = BaseCommand.option_list + (
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    14
        make_option('--fixtures', action='store_true', dest='infixtures', default=False,
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    15
            help='Only look in app.fixtures subdir'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    16
        make_option('--noscripts', action='store_true', dest='noscripts', default=False,
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    17
            help='Look in app.scripts subdir'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    18
        make_option('-s', '--silent', action='store_true', dest='silent', default=False,
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    19
            help='Run silently, do not show errors and tracebacks'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    20
        make_option('--no-traceback', action='store_true', dest='no_traceback', default=False,
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    21
            help='Do not show tracebacks'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    22
    )
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    23
    help = 'Runs a script in django context.'
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    24
    args = "script [script ...]"
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    25
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    26
    def handle(self, *scripts, **options):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    27
        from django.db.models import get_apps
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    28
        
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    29
        NOTICE = self.style.SQL_TABLE
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    30
        NOTICE2 = self.style.SQL_FIELD
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    31
        ERROR = self.style.ERROR_OUTPUT
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    32
        ERROR2 = self.style.NOTICE
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    33
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    34
        subdirs = []
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    35
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    36
        if not options.get('noscripts'):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    37
            subdirs.append('scripts')
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    38
        if options.get('infixtures'):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    39
            subdirs.append('fixtures')
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    40
        verbosity = int(options.get('verbosity', 1))
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    41
        show_traceback = options.get('traceback', True)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    42
        if show_traceback is None:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    43
            # XXX: traceback is set to None from Django ?
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    44
            show_traceback = True
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    45
        no_traceback = options.get('no_traceback', False)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    46
        if no_traceback:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    47
            show_traceback = False
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    48
        silent = options.get('silent', False)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    49
        if silent:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    50
            verbosity = 0
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    51
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    52
        if len(subdirs) < 1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    53
            print NOTICE("No subdirs to run left.")
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    54
            return
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    55
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    56
        if len(scripts) < 1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    57
            print ERROR("Script name required.")
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    58
            return
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    59
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    60
        def run_script(mod):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    61
            # TODO: add arguments to run
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    62
            try:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    63
                mod.run()
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    64
            except Exception, e:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    65
                if silent:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    66
                    return
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    67
                if verbosity > 0:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    68
                    print ERROR("Exception while running run() in '%s'" % mod.__name__)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    69
                if show_traceback:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    70
                    raise
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    71
        
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    72
        def my_import(mod):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    73
            if verbosity > 1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    74
                print NOTICE("Check for %s" % mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    75
            try:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    76
                t = __import__(mod, [], [], [" "])
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    77
                #if verbosity > 1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    78
                #    print NOTICE("Found script %s ..." % mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    79
                if hasattr(t, "run"):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    80
                    if verbosity > 1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    81
                        print NOTICE2("Found script '%s' ..." % mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    82
                    #if verbosity > 1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    83
                    #    print NOTICE("found run() in %s. executing..." % mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    84
                    return t
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    85
                else:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    86
                    if verbosity > 1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    87
                        print ERROR2("Find script '%s' but no run() function found." % mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    88
            except ImportError:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    89
                return False
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    90
        
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    91
        def find_modules_for_script(script):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    92
            """ find script module which contains 'run' attribute """
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    93
            modules = []
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    94
            # first look in apps
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    95
            for app in get_apps():
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    96
                app_name = app.__name__.split(".")[:-1] # + ['fixtures']
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    97
                for subdir in subdirs:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    98
                    mod = my_import(".".join(app_name + [subdir, script]))
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    99
                    if mod:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   100
                        modules.append(mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   101
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   102
            # try app.DIR.script import
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   103
            sa = script.split(".")
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   104
            for subdir in subdirs:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   105
                nn = ".".join(sa[:-1] + [subdir, sa[-1]])
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   106
                mod = my_import(nn)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   107
                if mod:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   108
                    modules.append(mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   109
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   110
            # try direct import
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   111
            if script.find(".") != -1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   112
                mod = my_import(script)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   113
                if mod:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   114
                    modules.append(mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   115
            
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   116
            return modules
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   117
        
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   118
        for script in scripts:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   119
            modules = find_modules_for_script(script)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   120
            if not modules:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   121
                if verbosity>0 and not silent:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   122
                    print ERROR("No module for script '%s' found" % script)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   123
            for mod in modules:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   124
                if verbosity>1:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   125
                    print NOTICE2("Running script '%s' ..." % mod.__name__)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   126
                run_script(mod)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   127
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   128
# Backwards compatibility for Django r9110
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   129
if not [opt for opt in Command.option_list if opt.dest=='verbosity']:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   130
    Command.option_list += (
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   131
        make_option('--verbosity', '-v', action="store", dest="verbosity",
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   132
                    default='1', type='choice', choices=['0', '1', '2'],
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   133
                    help="Verbosity level; 0=minimal output, 1=normal output, 2=all output"),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
   134
    )