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()
|
||||
|
||||
### 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:
|
||||
|
|
|
@ -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! "
|
||||
|
|
Loading…
Reference in a new issue