mirror of
https://github.com/veekun/pokedex.git
synced 2024-08-20 18:16:34 +00:00
Make load.py more idiomatic.
Also snuck an idiom into db/__init__.py.
This commit is contained in:
parent
b9a54c63d8
commit
bbb7cb9907
2 changed files with 10 additions and 10 deletions
|
@ -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:
|
||||||
|
|
|
@ -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! "
|
||||||
|
|
Loading…
Reference in a new issue