terom@59: #!/usr/bin/env python
terom@59: 
terom@59: """
terom@59:     Import inventory table from .csv
terom@59: """
terom@59: 
terom@59: from svv import application
terom@59: from svv import database as db
terom@59: from svv.items import Item
terom@59: 
terom@59: import csv
terom@59: 
terom@59: import optparse, logging
terom@59: 
terom@59: log = logging.getLogger('inventory-import')
terom@59: 
terom@59: def load_csv (path, encoding='utf-8') :
terom@59:     """
terom@59:         Load CSV data from given fs path, yielding the data as unicode strings.
terom@59:     """
terom@59: 
terom@59:     reader = csv.reader(open(path, 'r'))
terom@59: 
terom@59:     for row in reader :
terom@59:         # decode
terom@59:         yield [cell.decode(encoding) for cell in row]
terom@59: 
terom@59: def parse_csv (csv,
terom@59:     header_rows=3,
terom@59:     header_columns=[u"Sijainti", u"Artikkeli", u"Lkm", u"Selite"],
terom@59:     column_offset=0,
terom@59:     column_count=4,
terom@59:     normalize_qty=True,
terom@59: ) :
terom@59:     """
terom@59:         Parse item data from given csv:
terom@59:             
terom@59:             csv             - the csv.Reader object
terom@59: 
terom@59:             header_rows     - number of column header row
terom@59:             header_columns  - set of column headers to detect
terom@59:             column_offset   - number of columns to skip
terom@59:             column_count    - number of columns to use
terom@59:             normalize_qty   - items with qty=1 aren't countable
terom@59:         
terom@59:         Yields a series of
terom@59: 
terom@59:             (row_id, parent_row, name, qty, descr)
terom@59: 
terom@59:         tuples.
terom@59:     """
terom@59: 
terom@59:     # column def
terom@59:     column_def = ['location', 'name', 'qty', 'descr']
terom@59: 
terom@59:     # row_id's for loaded items
terom@59:     row_by_name = {}
terom@59: 
terom@59: 
terom@59:     for row, columns in enumerate(csv) :
terom@59:         # fix up row number
terom@59:         row += 1
terom@59: 
terom@59:         # fix up columns
terom@59:         columns = columns[column_offset:column_offset + column_count]
terom@59: 
terom@59:         # header
terom@59:         if row < header_rows :
terom@59:             # skip
terom@59:             continue
terom@59: 
terom@59:         elif row == header_rows :
terom@59:             # check header
terom@59:             if columns != header_columns :
terom@59:                 # fail
terom@59:                 raise Exception("Header mismatch, excepting %r, got %r on row %r" % (header_columns, columns, row))
terom@59: 
terom@59:             # ok
terom@59:             continue
terom@59: 
terom@59:         elif all(not col.strip() for col in columns) :
terom@59:             # empty row
terom@59:             continue
terom@59:         
terom@59:         elif len(columns) != len(column_def) :
terom@59:             # bad column count
terom@59:             raise Exception("Column count mismatch, expecting %r, got %r on row %r" % (column_def, columns, row))
terom@59:         
terom@59:         # load column
terom@59:         parent, name, qty, descr = columns
terom@59:         
terom@59:         # normalize
terom@59:         parent = parent.strip()
terom@59:         name = name.strip()
terom@59:         qty = qty.strip()
terom@59:         descr = descr.strip()
terom@59:         
terom@59:         # warn on whitespace
terom@59:         if [parent, name, qty, descr] != columns :
terom@59:             log.warning("Extra whitespace on row %d", row)
terom@59: 
terom@59:         # convert
terom@59:         try :
terom@59:             qty = int(qty) if qty else None
terom@59: 
terom@59:         except ValueError :
terom@59:             raise Exception("Invalid item quantity, got %r on row %r" % (qty, row))
terom@59: 
terom@59:         # normalize quantity
terom@59:         if qty == 0 :
terom@59:             # skip
terom@59:             log.warning("Skipping qty=0 item %r on row %d", name, row)
terom@59:             
terom@59:         elif qty == 1 and normalize_qty :
terom@59:             # one item is a single item
terom@59:             qty = None
terom@59: 
terom@59:         # map
terom@59:         if not parent :
terom@59:             parent = None
terom@59: 
terom@59:         elif parent not in row_by_name :
terom@59:             raise Exception("Invalid location, got %r on row %r" % (parent, row))
terom@59: 
terom@59:         else :
terom@59:             # lookup row_id
terom@59:             parent = row_by_name[parent]
terom@59:         
terom@59:         # debug
terom@59:         log.debug("Load %d: %s %s x%s -> %d", row, name, descr, qty, parent)
terom@59:         
terom@59:         # store
terom@59:         row_by_name[name] = row
terom@59: 
terom@59:         # yield
terom@59:         yield row, parent, name, qty, descr
terom@59: 
terom@59: def create_items (session, items) :
terom@59:     """
terom@59:         Create the given items and add them to the session
terom@59:     """
terom@59: 
terom@59:     # lookup parents
terom@59:     items_by_row = {}
terom@59: 
terom@59:     for row, parent, name, qty, descr in items :
terom@59:         # look up parent
terom@59:         if parent :
terom@59:             parent = items_by_row[parent]
terom@59: 
terom@59:         # build
terom@59:         item = Item(name, descr, qty, parent)
terom@59: 
terom@59:         # store
terom@59:         items_by_row[row] = item
terom@59: 
terom@59:         # add
terom@59:         session.add(item)
terom@59: 
terom@59:     # number of items added
terom@59:     return len(items_by_row)
terom@59: 
terom@59: def main () :
terom@59:     parser = optparse.OptionParser(usage="Usage: %prog [options] CSV")
terom@59: 
terom@59:     parser.add_option('-q', '--quiet', action='store_true', help='More output')
terom@59:     parser.add_option('-v', '--verbose', action='store_true', help='More output')
terom@59:     parser.add_option('-D', '--debug', action='store_true', help="Even more output (SQL queries)")
terom@59: 
terom@59:     parser.add_option('-d', '--database', help="Database connection URI", metavar='URL')
terom@59:     parser.add_option(      '--init-database', action='store_true', help="Initialize database (CREATE)")
terom@59: 
terom@59:     parser.add_option('-C', '--charset', help="CSV charset", metavar='CHARSET', default='utf-8')
terom@59: 
terom@59:     parser.add_option('--dry-run', action='store_true', help="Execute INSERTs but don't actually commit()")
terom@59: 
terom@59:     (options, args) = parser.parse_args()
terom@59: 
terom@59:     # CSV from args
terom@59:     csv_path, = args
terom@59:     csv_encoding = options.charset
terom@59:     
terom@59:     # defaults
terom@59:     level = logging.INFO
terom@59:     sql_echo = False
terom@59: 
terom@59:     if options.quiet :
terom@59:         level = logging.WARN
terom@59: 
terom@59:     if options.verbose :
terom@59:         level = logging.DEBUG
terom@59:     
terom@59:     if options.debug :
terom@59:         level = logging.DEBUG
terom@59:         sql_echo = True
terom@59: 
terom@59:     logging.basicConfig(format="[%(levelname)5s] %(funcName)25s : %(message)s", level=level)
terom@59: 
terom@59:     if sql_echo :
terom@59:         # echo SQL queries
terom@59:         logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
terom@59: 
terom@59:     # app state
terom@59:     app = application.Application(
terom@59:             options.database,
terom@59:     )
terom@59: 
terom@59:     # init?
terom@59:     if options.init_database :
terom@59:         app.create_tables()
terom@59: 
terom@59:     # load CSV data
terom@59:     csv = load_csv(csv_path, encoding=csv_encoding)
terom@59: 
terom@59:     # parse info
terom@59:     item_data = list(parse_csv(csv))
terom@59:     
terom@59:     # db session
terom@59:     session = app.session()
terom@59: 
terom@59:     # inserts
terom@59:     count = create_items(session, item_data)
terom@59:     
terom@59:     if options.dry_run :
terom@59:         log.info("Rolling back %d INSERTs for --dry-run...", count)
terom@59: 
terom@59:         # revert
terom@59:         session.rollback()
terom@59: 
terom@59:     else :
terom@59:         log.info("Committing %d new items..", count)
terom@59: 
terom@59:         # ok, commit
terom@59:         session.commit()
terom@59: 
terom@59: if __name__ == '__main__' :
terom@59:     main()
terom@59: