Make load.py more idiomatic.

Also snuck an idiom into db/__init__.py.
This commit is contained in:
a_magical_me 2011-04-03 04:26:45 -07:00
parent b9a54c63d8
commit bbb7cb9907
2 changed files with 10 additions and 10 deletions

View file

@ -21,7 +21,7 @@ def connect(uri=None, session_args={}, engine_args={}, engine_prefix=''):
uri = get_default_db_uri()
### Do some fixery for MySQL
if uri[0:5] == 'mysql':
if uri.startswith('mysql:'):
# MySQL uses latin1 for connections by default even if the server is
# otherwise oozing with utf8; charset fixes this
if 'charset' not in uri:

View file

@ -53,7 +53,7 @@ def _get_verbose_prints(verbose):
def print_start(thing):
# Truncate to 66 characters, leaving 10 characters for a success
# or failure message
truncated_thing = thing[0:66]
truncated_thing = thing[:66]
# Also, space-pad to keep the cursor in a known column
num_spaces = 66 - len(truncated_thing)
@ -212,12 +212,12 @@ def load(session, tables=[], directory=None, drop_tables=False, verbose=False, s
# them to the session last
# ASSUMPTION: Self-referential tables have a single PK called "id"
deferred_rows = [] # ( row referring to id, [foreign ids we need] )
seen_ids = {} # primary key we've seen => 1
seen_ids = set() # primary keys we've seen
# Fetch foreign key columns that point at this table, if any
self_ref_columns = []
for column in table_obj.c:
if any(_.references(table_obj) for _ in column.foreign_keys):
if any(x.references(table_obj) for x in column.foreign_keys):
self_ref_columns.append(column)
new_rows = []
@ -256,18 +256,18 @@ def load(session, tables=[], directory=None, drop_tables=False, verbose=False, s
# May need to stash this row and add it later if it refers to a
# later row in this table
if self_ref_columns:
foreign_ids = [row_data[_.name] for _ in self_ref_columns]
foreign_ids = [_ for _ in foreign_ids if _] # remove NULL ids
foreign_ids = set(row_data[x.name] for x in self_ref_columns)
foreign_ids.discard(None) # remove NULL ids
if not foreign_ids:
# NULL key. Remember this row and add as usual.
seen_ids[row_data['id']] = 1
seen_ids.add(row_data['id'])
elif all(_ in seen_ids for _ in foreign_ids):
elif foreign_ids.issubset(seen_ids):
# Non-NULL key we've already seen. Remember it and commit
# so we know the old row exists when we add the new one
insert_and_commit()
seen_ids[row_data['id']] = 1
seen_ids.add(row_data['id'])
else:
# Non-NULL future id. Save this and insert it later!
@ -286,7 +286,7 @@ def load(session, tables=[], directory=None, drop_tables=False, verbose=False, s
# Attempt to add any spare rows we've collected
for row_data, foreign_ids in deferred_rows:
if not all(_ in seen_ids for _ in foreign_ids):
if not foreign_ids.issubset(seen_ids):
# Could happen if row A refers to B which refers to C.
# This is ridiculous and doesn't happen in my data so far
raise ValueError("Too many levels of self-reference! "