继续上一章所讲，上一章我们最后面说道，虽然这个是很小的程序，但还有好几个要优化的地方。先复制一下老的view.py代码。
1 # coding:utf-8 2 from flask import Flask, request, jsonify 3 from model import User, db_session 4 import hashlib 5 import time 6 import redis 7 8 app = Flask(__name__) 9 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123') 10 11 12 @app.route('/') 13 def hello_world(): 14 return 'Hello World!' 15 16 17 @app.route('/login', methods=['POST']) 18 def login(): 19 phone_number = request.get_json().get('phone_number') 20 password = request.get_json().get('password') 21 user = User.query.filter_by(phone_number=phone_number).first() 22 if not user: 23 return jsonify({'code': 0, 'message': '没有此用户'}) 24 25 if user.password != password: 26 return jsonify({'code': 0, 'message': '密码错误'}) 27 28 m = hashlib.md5() 29 m.update(phone_number) 30 m.update(password) 31 m.update(str(int(time.time()))) 32 token = m.hexdigest() 33 34 redis_store.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1}) 35 redis_store.set('token:%s' % token, user.phone_number) 36 redis_store.expire('token:%s' % token, 3600*24*30) 37 38 return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token}) 39 40 41 @app.route('/user') 42 def user(): 43 token = request.headers.get('token') 44 if not token: 45 return jsonify({'code': 0, 'message': '需要验证'}) 46 phone_number = redis_store.get('token:%s' % token) 47 if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'): 48 return jsonify({'code': 2, 'message': '验证信息错误'}) 49 50 nickname = redis_store.hget('user:%s' % phone_number, 'nickname') 51 return jsonify({'code': 1, 'nickname': nickname, 'phone_number': phone_number}) 52 53 54 @app.route('/logout') 55 def logout(): 56 token = request.headers.get('token') 57 if not token: 58 return jsonify({'code': 0, 'message': '需要验证'}) 59 phone_number = redis_store.get('token:%s' % token) 60 if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'): 61 return jsonify({'code': 2, 'message': '验证信息错误'}) 62 63 redis_store.delete('token:%s' % token) 64 redis_store.hmset('user:%s' % phone_number, {'app_online': 0}) 65 return jsonify({'code': 1, 'message': '成功注销'}) 66 67 68 @app.teardown_request 69 def handle_teardown_request(exception): 70 db_session.remove() 71 72 if __name__ == '__main__': 73 app.run(debug=True, host='0.0.0.0', port=5001)
其中验证token的方法，已经重叠了，python教我们，永远不要重复自己的代码，这是很丑陋的行为。今天我们把它换成一个装饰器，然后再把redis调整一下，看看代码会不会简洁很多。
1 # coding:utf-8 2 from flask import Flask, request, jsonify 3 from model import User, db_session 4 import hashlib 5 import time 6 import redis 7 from functools import wraps 8 9 app = Flask(__name__) 10 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123') 11 12 13 def login_check(f): 14 @wraps(f) 15 def decorator(*args, **kwargs): 16 token = request.headers.get('token') 17 if not token: 18 return jsonify({'code': 0, 'message': '需要验证'}) 19 20 phone_number = redis_store.get('token:%s' % token) 21 if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'): 22 return jsonify({'code': 2, 'message': '验证信息错误'}) 23 24 return f(*args, **kwargs) 25 return decorator 26 27 28 @app.route('/login', methods=['POST']) 29 def login(): 30 phone_number = request.get_json().get('phone_number') 31 password = request.get_json().get('password') 32 user = User.query.filter_by(phone_number=phone_number).first() 33 if not user: 34 return jsonify({'code': 0, 'message': '没有此用户'}) 35 36 if user.password != password: 37 return jsonify({'code': 0, 'message': '密码错误'}) 38 39 m = hashlib.md5() 40 m.update(phone_number) 41 m.update(password) 42 m.update(str(int(time.time()))) 43 token = m.hexdigest() 44 45 pipeline = redis_store.pipeline() 46 pipeline.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1}) 47 pipeline.set('token:%s' % token, user.phone_number) 48 pipeline.expire('token:%s' % token, 3600*24*30) 49 pipeline.execute() 50 51 return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token}) 52 53 54 @app.route('/user') 55 @login_check 56 def user(): 57 token = request.headers.get('token') 58 phone_number = redis_store.get('token:%s' % token) 59 60 nickname = redis_store.hget('user:%s' % phone_number, 'nickname') 61 return jsonify({'code': 1, 'nickname': nickname, 'phone_number': phone_number}) 62 63 64 @app.route('/logout') 65 @login_check 66 def logout(): 67 token = request.headers.get('token') 68 phone_number = redis_store.get('token:%s' % token) 69 70 pipeline = redis_store.pipeline() 71 pipeline.delete('token:%s' % token) 72 pipeline.hmset('user:%s' % phone_number, {'app_online': 0}) 73 pipeline.execute() 74 return jsonify({'code': 1, 'message': '成功注销'}) 75 76 77 @app.teardown_request 78 def handle_teardown_request(exception): 79 db_session.remove() 80 81 if __name__ == '__main__': 82 app.run(debug=True, host='0.0.0.0', port=5001)
加了一个装饰器，是不是简洁了很多？每次需要验证的时候，只需要一个login_check就可以了，这样就变得非常简洁，而且脉络清晰。redis也改成了管道执行，pipeline，防止执行到一半，被掐断。
可是，可是，我还是觉得不简洁，看user, logout的代码中重复的地方。
token = request.headers.get('token') phone_number = redis_store.get('token:%s' % token)
每次都有这两句，要是将来还有其他值怎么办？上面不刚说，永远不要重复自己的代码吗？
好，我们再写一个函数，看下面代码
1 # coding:utf-8 2 from flask import Flask, request, jsonify, g 3 from model import User, db_session 4 import hashlib 5 import time 6 import redis 7 from functools import wraps 8 9 app = Flask(__name__) 10 redis_store = redis.Redis(host='localhost', port=6380, db=4, password='dahai123') 11 12 13 def login_check(f): 14 @wraps(f) 15 def decorator(*args, **kwargs): 16 token = request.headers.get('token') 17 if not token: 18 return jsonify({'code': 0, 'message': '需要验证'}) 19 20 phone_number = redis_store.get('token:%s' % token) 21 if not phone_number or token != redis_store.hget('user:%s' % phone_number, 'token'): 22 return jsonify({'code': 2, 'message': '验证信息错误'}) 23 24 return f(*args, **kwargs) 25 return decorator 26 27 28 @app.before_request 29 def before_request(): 30 token = request.headers.get('token') 31 phone_number = redis_store.get('token:%s' % token) 32 if phone_number: 33 g.current_user = User.query.filter_by(phone_number=phone_number).first() 34 g.token = token 35 return 36 37 38 @app.route('/login', methods=['POST']) 39 def login(): 40 phone_number = request.get_json().get('phone_number') 41 password = request.get_json().get('password') 42 user = User.query.filter_by(phone_number=phone_number).first() 43 if not user: 44 return jsonify({'code': 0, 'message': '没有此用户'}) 45 46 if user.password != password: 47 return jsonify({'code': 0, 'message': '密码错误'}) 48 49 m = hashlib.md5() 50 m.update(phone_number) 51 m.update(password) 52 m.update(str(int(time.time()))) 53 token = m.hexdigest() 54 55 pipeline = redis_store.pipeline() 56 pipeline.hmset('user:%s' % user.phone_number, {'token': token, 'nickname': user.nickname, 'app_online': 1}) 57 pipeline.set('token:%s' % token, user.phone_number) 58 pipeline.expire('token:%s' % token, 3600*24*30) 59 pipeline.execute() 60 61 return jsonify({'code': 1, 'message': '成功登录', 'nickname': user.nickname, 'token': token}) 62 63 64 @app.route('/user') 65 @login_check 66 def user(): 67 user = g.current_user 68 69 nickname = redis_store.hget('user:%s' % user.phone_number, 'nickname') 70 return jsonify({'code': 1, 'nickname': nickname, 'phone_number': user.phone_number}) 71 72 73 @app.route('/logout') 74 @login_check 75 def logout(): 76 user = g.current_user 77 78 pipeline = redis_store.pipeline() 79 pipeline.delete('token:%s' % g.token) 80 pipeline.hmset('user:%s' % user.phone_number, {'app_online': 0}) 81 pipeline.execute() 82 return jsonify({'code': 1, 'message': '成功注销'}) 83 84 85 @app.teardown_request 86 def handle_teardown_request(exception): 87 db_session.remove() 88 89 if __name__ == '__main__': 90 app.run(debug=True, host='0.0.0.0', port=5001)
我们在代码中加了一个before_request，这个函数就是在每个request发起的时候，如果已经验证了，我们把当前g.current_user和g.token设置一下，这样每次需要获取当前用户的时候，直接找g.current_user就可以了，是不是简单了太多太多？好了，今天到此为止，下一章，我们讲怎么利用alembic修改数据库。