1 #!/usr/bin/env python

     2 

     3 """

     4     Import inventory table from .csv

     5 """

     6 

     7 from svv import application

     8 from svv import database as db

     9 from svv.items import Item

    10 

    11 import csv

    12 

    13 import optparse, logging

    14 

    15 log = logging.getLogger('inventory-import')

    16 

    17 def load_csv (path, encoding='utf-8') :

    18     """

    19         Load CSV data from given fs path, yielding the data as unicode strings.

    20     """

    21 

    22     reader = csv.reader(open(path, 'r'))

    23 

    24     for row in reader :

    25         # decode

    26         yield [cell.decode(encoding) for cell in row]

    27 

    28 def parse_csv (csv,

    29     header_rows=3,

    30     header_columns=[u"Sijainti", u"Artikkeli", u"Lkm", u"Selite"],

    31     column_offset=0,

    32     column_count=4,

    33     normalize_qty=True,

    34 ) :

    35     """

    36         Parse item data from given csv:

    37             

    38             csv             - the csv.Reader object

    39 

    40             header_rows     - number of column header row

    41             header_columns  - set of column headers to detect

    42             column_offset   - number of columns to skip

    43             column_count    - number of columns to use

    44             normalize_qty   - items with qty=1 aren't countable

    45         

    46         Yields a series of

    47 

    48             (row_id, parent_row, name, qty, descr)

    49 

    50         tuples.

    51     """

    52 

    53     # column def

    54     column_def = ['location', 'name', 'qty', 'descr']

    55 

    56     # row_id's for loaded items

    57     row_by_name = {}

    58 

    59 

    60     for row, columns in enumerate(csv) :

    61         # fix up row number

    62         row += 1

    63 

    64         # fix up columns

    65         columns = columns[column_offset:column_offset + column_count]

    66 

    67         # header

    68         if row < header_rows :

    69             # skip

    70             continue

    71 

    72         elif row == header_rows :

    73             # check header

    74             if columns != header_columns :

    75                 # fail

    76                 raise Exception("Header mismatch, excepting %r, got %r on row %r" % (header_columns, columns, row))

    77 

    78             # ok

    79             continue

    80 

    81         elif all(not col.strip() for col in columns) :

    82             # empty row

    83             continue

    84         

    85         elif len(columns) != len(column_def) :

    86             # bad column count

    87             raise Exception("Column count mismatch, expecting %r, got %r on row %r" % (column_def, columns, row))

    88         

    89         # load column

    90         parent, name, qty, descr = columns

    91         

    92         # normalize

    93         parent = parent.strip()

    94         name = name.strip()

    95         qty = qty.strip()

    96         descr = descr.strip()

    97         

    98         # warn on whitespace

    99         if [parent, name, qty, descr] != columns :

   100             log.warning("Extra whitespace on row %d", row)

   101 

   102         # convert

   103         try :

   104             qty = int(qty) if qty else None

   105 

   106         except ValueError :

   107             raise Exception("Invalid item quantity, got %r on row %r" % (qty, row))

   108 

   109         # normalize quantity

   110         if qty == 0 :

   111             # skip

   112             log.warning("Skipping qty=0 item %r on row %d", name, row)

   113             

   114         elif qty == 1 and normalize_qty :

   115             # one item is a single item

   116             qty = None

   117 

   118         # map

   119         if not parent :

   120             parent = None

   121 

   122         elif parent not in row_by_name :

   123             raise Exception("Invalid location, got %r on row %r" % (parent, row))

   124 

   125         else :

   126             # lookup row_id

   127             parent = row_by_name[parent]

   128         

   129         # debug

   130         log.debug("Load %d: %s %s x%s -> %d", row, name, descr, qty, parent)

   131         

   132         # store

   133         row_by_name[name] = row

   134 

   135         # yield

   136         yield row, parent, name, qty, descr

   137 

   138 def create_items (session, items) :

   139     """

   140         Create the given items and add them to the session

   141     """

   142 

   143     # lookup parents

   144     items_by_row = {}

   145 

   146     for row, parent, name, qty, descr in items :

   147         # look up parent

   148         if parent :

   149             parent = items_by_row[parent]

   150 

   151         # build

   152         item = Item(name, descr, qty, parent)

   153 

   154         # store

   155         items_by_row[row] = item

   156 

   157         # add

   158         session.add(item)

   159 

   160     # number of items added

   161     return len(items_by_row)

   162 

   163 def main () :

   164     parser = optparse.OptionParser(usage="Usage: %prog [options] CSV")

   165 

   166     parser.add_option('-q', '--quiet', action='store_true', help='More output')

   167     parser.add_option('-v', '--verbose', action='store_true', help='More output')

   168     parser.add_option('-D', '--debug', action='store_true', help="Even more output (SQL queries)")

   169 

   170     parser.add_option('-d', '--database', help="Database connection URI", metavar='URL')

   171     parser.add_option(      '--init-database', action='store_true', help="Initialize database (CREATE)")

   172 

   173     parser.add_option('-C', '--charset', help="CSV charset", metavar='CHARSET', default='utf-8')

   174 

   175     parser.add_option('--dry-run', action='store_true', help="Execute INSERTs but don't actually commit()")

   176 

   177     (options, args) = parser.parse_args()

   178 

   179     # CSV from args

   180     csv_path, = args

   181     csv_encoding = options.charset

   182     

   183     # defaults

   184     level = logging.INFO

   185     sql_echo = False

   186 

   187     if options.quiet :

   188         level = logging.WARN

   189 

   190     if options.verbose :

   191         level = logging.DEBUG

   192     

   193     if options.debug :

   194         level = logging.DEBUG

   195         sql_echo = True

   196 

   197     logging.basicConfig(format="[%(levelname)5s] %(funcName)25s : %(message)s", level=level)

   198 

   199     if sql_echo :

   200         # echo SQL queries

   201         logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)

   202 

   203     # app state

   204     app = application.Application(

   205             options.database,

   206     )

   207 

   208     # init?

   209     if options.init_database :

   210         app.create_tables()

   211 

   212     # load CSV data

   213     csv = load_csv(csv_path, encoding=csv_encoding)

   214 

   215     # parse info

   216     item_data = list(parse_csv(csv))

   217     

   218     # db session

   219     session = app.session()

   220 

   221     # inserts

   222     count = create_items(session, item_data)

   223     

   224     if options.dry_run :

   225         log.info("Rolling back %d INSERTs for --dry-run...", count)

   226 

   227         # revert

   228         session.rollback()

   229 

   230     else :

   231         log.info("Committing %d new items..", count)

   232 

   233         # ok, commit

   234         session.commit()

   235 

   236 if __name__ == '__main__' :

   237     main()

   238