diff --git a/pokedex/cli/__init__.py b/pokedex/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pokedex/cli/search.py b/pokedex/cli/search.py new file mode 100644 index 0000000..cef0458 --- /dev/null +++ b/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) diff --git a/pokedex/main.py b/pokedex/main.py index 9b381ad..5758df1 100644 --- a/pokedex/main.py +++ b/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 diff --git a/pokedex/search.py b/pokedex/search.py new file mode 100644 index 0000000..fa5a04a --- /dev/null +++ b/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[-_a-zA-Z0-9]+): )? + (?P + (?: + [^\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()