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() uri = get_default_db_uri()
### Do some fixery for MySQL ### 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 # MySQL uses latin1 for connections by default even if the server is
# otherwise oozing with utf8; charset fixes this # otherwise oozing with utf8; charset fixes this
if 'charset' not in uri: if 'charset' not in uri:

View file

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