Stub out a CLI search interface, which can also work with JSON and strings

pokedex/cli/search.py

@@ -0,0 +1,22 @@
+from pokedex.search import search
+
+
+def configure_parser(parser):
+    parser.set_defaults(func=command_search)
+
+    parser.add_argument('--name', default=None)
+
+    parser.add_argument('--attack', '--atk', dest='attack', default=None)
+    parser.add_argument('--defense', '--def', dest='defense', default=None)
+    parser.add_argument('--special-attack', '--spatk', dest='special-attack', default=None)
+    parser.add_argument('--special-defense', '--spdef', dest='special-defense', default=None)
+    parser.add_argument('--speed', dest='speed', default=None)
+    parser.add_argument('--hp', dest='hp', default=None)
+
+
+def command_search(parser, args):
+    from pokedex.main import get_session
+    session = get_session(args)
+    results = search(session, **vars(args))
+    for result in results:
+        print(result.name)

pokedex/main.py

@@ -5,6 +5,7 @@ import argparse
 import os
 import sys
 
+import pokedex.cli.search
 import pokedex.db
 import pokedex.db.load
 import pokedex.db.tables
@@ -63,6 +64,9 @@ def create_parser():
     cmd_lookup.set_defaults(func=command_lookup)
     cmd_lookup.add_argument('criteria', nargs='+')
 
+    cmd_search = cmds.add_parser('search', help=u'Find things by various criteria')
+    pokedex.cli.search.configure_parser(cmd_search)
+
     cmd_load = cmds.add_parser('load', help=u'Load Pokédex data into a database from CSV files')
     cmd_load.set_defaults(func=command_load, verbose=True)
     # TODO get the actual default here

pokedex/search.py

@@ -0,0 +1,65 @@
+import re
+
+from sqlalchemy import func
+from sqlalchemy.orm import joinedload
+
+import pokedex.db.tables as t
+
+
+def _parse_range(value):
+    v = int(value)
+    return lambda x: x == v
+
+
+CRITERION_RX = re.compile(r"""
+    \s*
+    (?: (?P<field>[-_a-zA-Z0-9]+): )?
+    (?P<pattern>
+        (?:
+            [^\s"]+?
+        )+
+    )
+""", re.VERBOSE)
+def parse_search_string(string):
+    """Parses a search string!"""
+    criteria = {}
+    for match in CRITERION_RX.finditer(string):
+        # TODO what if there are several of the same match!
+        # TODO the cli needs to do append too
+        field = match.group('field') or '*'
+        criteria[field] = match.group('pattern')
+    return criteria
+
+
+def search(session, **criteria):
+    query = (
+        session.query(t.Pokemon)
+        .options(
+            joinedload(t.Pokemon.species)
+        )
+    )
+
+    stat_query = (
+        session.query(t.PokemonStat.pokemon_id)
+        .join(t.PokemonStat.stat)
+    )
+    do_stat = False
+
+    if criteria.get('name') is not None:
+        query = query.filter(t.Pokemon.species.has(func.lower(t.PokemonSpecies.name) == criteria['name'].lower()))
+
+    for stat_ident in (u'attack', u'defense', u'special-attack', u'special-defense', u'speed', u'hp'):
+        criterion = criteria.get(stat_ident)
+        if criterion is None:
+            continue
+
+        do_stat = True
+        stat_query = stat_query.filter(
+            (t.Stat.identifier == stat_ident)
+            & _parse_range(criterion)(t.PokemonStat.base_stat)
+        )
+
+    if do_stat:
+        query = query.filter(t.Pokemon.id.in_(stat_query.subquery()))
+
+    return query.all()