colly/commands/__init__.py

   1 import sys
   2 from pkgutil import walk_packages
   3 import logging
   4 import optparse
   5
   6 from colly.optbase import parser
   7 from colly.core import Collate
   8
   9 __all__ = ['command_dict', 'Command', 'load_command',
  10            'load_all_commands', 'command_names']
  11
  12 command_dict = {}
  13
  14 ''' Base command
  15 '''
  16
  17 class Command(object):
  18     name = None
  19     usage = None
  20
  21     def __init__(self):
  22         assert self.name
  23         self.parser = optparse.OptionParser(
  24             usage=self.usage,
  25             prog='%s %s' % ("colly", self.name),
  26             version=parser.version)
  27         for option in parser.option_list:
  28             if not option.dest or option.dest == 'help':
  29                 # -h, --version, etc
  30                 continue
  31             self.parser.add_option(option)
  32         command_dict[self.name] = self
  33
  34     def main(self, initial_options, args):
  35         options, args = self.parser.parse_args(args)
  36         self.run(options, args)
  37
  38
  39 ''' CSV base command
  40     --
  41     @ Parses common options to initialize colly's core(.py) object
  42 '''
  43
  44 class CsvCommand(Command):
  45     usage = "%prog FILE(S) [OPTIONS]"
  46     objects = []
  47
  48     def __init__(self):
  49         super(CsvCommand, self).__init__()
  50
  51         self.parser.add_option('-H', '--headings',
  52             dest='headings',
  53             action='callback',
  54             callback=split_callback,
  55             default=[],
  56             type='str',
  57             help='Comma separated headings for CSV columns')
  58
  59         self.parser.add_option('-i', '--index',
  60             dest='index',
  61             action='callback',
  62             callback=split_callback,
  63             default=[],
  64             type='str',
  65             help='Define index column on CSV, shortcut to using -H pk,etc.')
  66
  67         '''
  68         TODO:
  69         self.parser.add_option('-z', '--lazy',
  70             dest='lazy',
  71             action='store_true',
  72             help='Will make assumptions rather than throw errors')
  73
  74         self.parser.add_option('-i', '--index',
  75             dest='pk',
  76             action='store',
  77             help='Column to set as the index (pk)')
  78
  79         self.parser.add_option('-f', '--format',
  80             dest='format',
  81             action='store',
  82             help='Output format [default: JSON]')
  83
  84         self.parser.add_option('-c', '--compress',
  85             dest='compress',
  86             action='store_true',
  87             help='Column to set as the index (pk)')
  88         '''
  89
  90     def main(self, initial_options, args): #:! overwrites Command.main
  91         options, args = self.parser.parse_args(args)
  92
  93         if options.headings and options.index:
  94             self.parser.error('Options HEADINGS and INDEX are mutually exclusive')
  95
  96         n = 0
  97         for csv_file in args:
  98             kwargs = {}
  99
 100             if options.headings:
 101                 try:
 102                     kwargs['headings'] = options.headings[n]
 103                 except IndexError:
 104                     break
 105             elif options.index:
 106                  try:
 107                      kwargs['headings'] = [] # leave out index for now
 108                  except IndexError:
 109                      break
 110             else:
 111                  pass
 112
 113             self.objects.append(Collate(csv_file, **kwargs))
 114             n += 1
 115
 116         self.run(options, args)
 117
 118 # Utils
 119
 120 def split_callback(option, opt, value, parser):
 121     ''' Callback, extra option parsing
 122     '''
 123     value_list = map(lambda v: v.split(','), value.split(':'))
 124     setattr(parser.values, option.dest, value_list)
 125
 126
 127 def load_command(name):
 128     full_name = 'colly.commands.%s' % name
 129     if full_name in sys.modules:
 130         return
 131     try:
 132         __import__(full_name)
 133     except ImportError:
 134         pass
 135
 136 def load_all_commands():
 137     for name in command_names():
 138         load_command(name)
 139
 140 def command_names():
 141     names = set((pkg[1] for pkg in walk_packages(path=__path__)))
 142     return list(names)