Move the rest of pokedex.util.get to pokedex.db.util

pokedex/db/util.py

@@ -0,0 +1,123 @@
+"""Helpers for common ways to work with pokedex queries
+
+These include identifier- and name-based lookup, filtering out base forms
+of pokemon, and filtering/ordering by name.
+"""
+
+from sqlalchemy.orm import aliased
+
+from pokedex.db import tables
+
+### Getter
+
+def get(session, table, identifier=None, name=None, id=None,
+        form_identifier=None, form_name=None, language=None, is_pokemon=None):
+    """Get one object from the database.
+
+    session: The session to use (from pokedex.db.connect())
+    table: The table to select from (such as pokedex.db.tables.Move)
+
+    identifier: Identifier of the object
+    name: The name of the object
+    id: The ID number of the object
+    form_identifier: For pokemon, identifier of the form
+    form_name: For pokemon, name of the form
+
+    language: A Language to use for name and form_name
+    is_pokemon: If true, specifies that the table should be treated as a
+        pokemon table (handling forms specially). If None and table is the
+        (unaliased) Pokemon, it is set to True. Otherwise, the pokemon forms
+        aren't handled.
+
+    All conditions must match, so it's not a good idea to specify more than one
+    of identifier/name/id at once.
+
+    If zero or more than one objects matching the criteria are found, the
+    appropriate SQLAlchemy exception is raised.
+    Exception: for pokemon, selects the form base unless form_* is given.
+    """
+
+    if is_pokemon is None:
+        is_pokemon = (table is tables.Pokemon)
+
+    query = session.query(table)
+
+    if identifier is not None:
+        query = query.filter_by(identifier=identifier)
+
+    if name is not None:
+        query = filter_name(query, table, name, language)
+
+    if id is not None:
+        query = query.filter_by(id=id)
+
+    if form_identifier is not None or form_name is not None:
+        if is_pokemon:
+            query = query.join(table.unique_form)
+            if form_identifier is not None:
+                query = query.filter(tables.PokemonForm.identifier ==
+                        form_identifier)
+            if form_name is not None:
+                query = filter_name(query, table, form_name, language)
+        else:
+            raise ValueError(
+                "form_identifier and form_name only make sense for pokemon")
+    elif is_pokemon:
+        query = filter_base_forms(query)
+
+    return query.one()
+
+### Helpers
+
+def filter_name(query, table, name, language):
+    """Filter a query by name, return the resulting query
+
+    query: The query to filter
+    table: The table of named objects
+    name: The name to look for. May be a tuple of alternatives.
+    language: The language for "name", or None for the session default
+    """
+    if language is None:
+        query = query.filter(table.name == name)
+    else:
+        names_table = table.names_table
+        query = query.join(names_table)
+        query = query.filter(names_table.foreign_id == table.id)
+        query = query.filter(names_table.local_language_id == language.id)
+        if isinstance(name, tuple):
+            query = query.filter(names_table.name in name)
+        else:
+            query = query.filter(names_table.name == name)
+    return query
+
+def filter_base_forms(query):
+    """Filter only base forms of pokemon, and return the resulting query
+    """
+    query = query.filter(tables.Pokemon.forms.any())
+    return query
+
+def order_by_name(query, table, language=None, *extra_languages):
+    """Order a query by name.
+
+    query: The query to order
+    table: Table of the named objects
+    language: The language to order names by. If None, use the
+        connection default.
+    extra_languages: Extra languages to order by, should the translations for
+        `language` be incomplete (or ambiguous).
+
+    Uses the identifier as a fallback ordering.
+    """
+    if language is None:
+        query = query.outerjoin(table.names_local)
+        query = query.order_by(table.names_table.name)
+    else:
+        extra_languages = (language, ) + extra_languages
+    for language in extra_languages:
+        names_table = aliased(table.names_table)
+        query = query.outerjoin(names_table)
+        query = query.filter(names_table.foreign_id == table.id)
+        query = query.filter(names_table.local_language_id == language.id)
+        query = query.order_by(names_table.name)
+    query = query.order_by(table.identifier)
+    return query