diff --git a/oyster/web.py b/oyster/web.py index 12c0ffb..5071fa7 100644 --- a/oyster/web.py +++ b/oyster/web.py @@ -1,5 +1,6 @@ import json import datetime +import functools import flask import pymongo.objectid @@ -17,37 +18,54 @@ class JSONEncoder(json.JSONEncoder): return super(JSONEncoder, self).default(obj) +def api_wrapper(template=None): + def wrapper(func): + @functools.wraps(func) + def newfunc(*args, **kwargs): + data = func(*args, **kwargs) + if 'json' in flask.request.args or not template: + return json.dumps(data, cls=JSONEncoder) + else: + return flask.render_template(template, **data) + + return newfunc + return wrapper + + app = flask.Flask('oyster') client = Client() + @app.route('/status/') +@api_wrapper() def doc_list(): status = { 'queue_size': app.work_queue.qsize(), 'tracking': client.db.tracked.count(), 'need_update': client.get_update_queue_size(), } - return json.dumps(status) + return status @app.route('/log/') +@api_wrapper('logs.html') def log_view(): offset = int(flask.request.args.get('offset', 0)) size = 100 prev_offset = max(offset - size, 0) next_offset = offset + size logs = client.db.logs.find().sort('$natural', -1).skip(offset).limit(size) - return flask.render_template('logs.html', logs=logs, - prev_offset=prev_offset, - next_offset=next_offset, - offset=offset) + return dict(logs=logs, prev_offset=prev_offset, + next_offset=next_offset, offset=offset) @app.route('/tracked/') +@api_wrapper() def tracked(): tracked = list(client.db.tracked.find()) return json.dumps(tracked, cls=JSONEncoder) + @app.route('/tracked/') def tracked_view(url): doc = client.db.tracked.find_one({'url': url})