__init__.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716
  1. from __future__ import absolute_import
  2. from functools import wraps, partial
  3. from flask import request, url_for, current_app
  4. from flask import abort as original_flask_abort
  5. from flask import make_response as original_flask_make_response
  6. from flask.views import MethodView
  7. from flask.signals import got_request_exception
  8. from werkzeug.datastructures import Headers
  9. from werkzeug.exceptions import HTTPException, MethodNotAllowed, NotFound, NotAcceptable, InternalServerError
  10. from werkzeug.wrappers import Response as ResponseBase
  11. from flask_restful.utils import http_status_message, unpack, OrderedDict
  12. from flask_restful.representations.json import output_json
  13. import sys
  14. from types import MethodType
  15. import operator
  16. try:
  17. from collections.abc import Mapping
  18. except ImportError:
  19. from collections import Mapping
  20. __all__ = ('Api', 'Resource', 'marshal', 'marshal_with', 'marshal_with_field', 'abort')
  21. def abort(http_status_code, **kwargs):
  22. """Raise a HTTPException for the given http_status_code. Attach any keyword
  23. arguments to the exception for later processing.
  24. """
  25. #noinspection PyUnresolvedReferences
  26. try:
  27. original_flask_abort(http_status_code)
  28. except HTTPException as e:
  29. if len(kwargs):
  30. e.data = kwargs
  31. raise
  32. DEFAULT_REPRESENTATIONS = [('application/json', output_json)]
  33. class Api(object):
  34. """
  35. The main entry point for the application.
  36. You need to initialize it with a Flask Application: ::
  37. >>> app = Flask(__name__)
  38. >>> api = restful.Api(app)
  39. Alternatively, you can use :meth:`init_app` to set the Flask application
  40. after it has been constructed.
  41. :param app: the Flask application object
  42. :type app: flask.Flask or flask.Blueprint
  43. :param prefix: Prefix all routes with a value, eg v1 or 2010-04-01
  44. :type prefix: str
  45. :param default_mediatype: The default media type to return
  46. :type default_mediatype: str
  47. :param decorators: Decorators to attach to every resource
  48. :type decorators: list
  49. :param catch_all_404s: Use :meth:`handle_error`
  50. to handle 404 errors throughout your app
  51. :param serve_challenge_on_401: Whether to serve a challenge response to
  52. clients on receiving 401. This usually leads to a username/password
  53. popup in web browsers.
  54. :param url_part_order: A string that controls the order that the pieces
  55. of the url are concatenated when the full url is constructed. 'b'
  56. is the blueprint (or blueprint registration) prefix, 'a' is the api
  57. prefix, and 'e' is the path component the endpoint is added with
  58. :type catch_all_404s: bool
  59. :param errors: A dictionary to define a custom response for each
  60. exception or error raised during a request
  61. :type errors: dict
  62. """
  63. def __init__(self, app=None, prefix='',
  64. default_mediatype='application/json', decorators=None,
  65. catch_all_404s=False, serve_challenge_on_401=False,
  66. url_part_order='bae', errors=None):
  67. self.representations = OrderedDict(DEFAULT_REPRESENTATIONS)
  68. self.urls = {}
  69. self.prefix = prefix
  70. self.default_mediatype = default_mediatype
  71. self.decorators = decorators if decorators else []
  72. self.catch_all_404s = catch_all_404s
  73. self.serve_challenge_on_401 = serve_challenge_on_401
  74. self.url_part_order = url_part_order
  75. self.errors = errors or {}
  76. self.blueprint_setup = None
  77. self.endpoints = set()
  78. self.resources = []
  79. self.app = None
  80. self.blueprint = None
  81. if app is not None:
  82. self.app = app
  83. self.init_app(app)
  84. def init_app(self, app):
  85. """Initialize this class with the given :class:`flask.Flask`
  86. application or :class:`flask.Blueprint` object.
  87. :param app: the Flask application or blueprint object
  88. :type app: flask.Flask
  89. :type app: flask.Blueprint
  90. Examples::
  91. api = Api()
  92. api.add_resource(...)
  93. api.init_app(app)
  94. """
  95. # If app is a blueprint, defer the initialization
  96. try:
  97. app.record(self._deferred_blueprint_init)
  98. # Flask.Blueprint has a 'record' attribute, Flask.Api does not
  99. except AttributeError:
  100. self._init_app(app)
  101. else:
  102. self.blueprint = app
  103. def _complete_url(self, url_part, registration_prefix):
  104. """This method is used to defer the construction of the final url in
  105. the case that the Api is created with a Blueprint.
  106. :param url_part: The part of the url the endpoint is registered with
  107. :param registration_prefix: The part of the url contributed by the
  108. blueprint. Generally speaking, BlueprintSetupState.url_prefix
  109. """
  110. parts = {
  111. 'b': registration_prefix,
  112. 'a': self.prefix,
  113. 'e': url_part
  114. }
  115. return ''.join(parts[key] for key in self.url_part_order if parts[key])
  116. @staticmethod
  117. def _blueprint_setup_add_url_rule_patch(blueprint_setup, rule, endpoint=None, view_func=None, **options):
  118. """Method used to patch BlueprintSetupState.add_url_rule for setup
  119. state instance corresponding to this Api instance. Exists primarily
  120. to enable _complete_url's function.
  121. :param blueprint_setup: The BlueprintSetupState instance (self)
  122. :param rule: A string or callable that takes a string and returns a
  123. string(_complete_url) that is the url rule for the endpoint
  124. being registered
  125. :param endpoint: See BlueprintSetupState.add_url_rule
  126. :param view_func: See BlueprintSetupState.add_url_rule
  127. :param **options: See BlueprintSetupState.add_url_rule
  128. """
  129. if callable(rule):
  130. rule = rule(blueprint_setup.url_prefix)
  131. elif blueprint_setup.url_prefix:
  132. rule = blueprint_setup.url_prefix + rule
  133. options.setdefault('subdomain', blueprint_setup.subdomain)
  134. if endpoint is None:
  135. endpoint = view_func.__name__
  136. defaults = blueprint_setup.url_defaults
  137. if 'defaults' in options:
  138. defaults = dict(defaults, **options.pop('defaults'))
  139. blueprint_setup.app.add_url_rule(rule, '%s.%s' % (blueprint_setup.blueprint.name, endpoint),
  140. view_func, defaults=defaults, **options)
  141. def _deferred_blueprint_init(self, setup_state):
  142. """Synchronize prefix between blueprint/api and registration options, then
  143. perform initialization with setup_state.app :class:`flask.Flask` object.
  144. When a :class:`flask_restful.Api` object is initialized with a blueprint,
  145. this method is recorded on the blueprint to be run when the blueprint is later
  146. registered to a :class:`flask.Flask` object. This method also monkeypatches
  147. BlueprintSetupState.add_url_rule with _blueprint_setup_add_url_rule_patch.
  148. :param setup_state: The setup state object passed to deferred functions
  149. during blueprint registration
  150. :type setup_state: flask.blueprints.BlueprintSetupState
  151. """
  152. self.blueprint_setup = setup_state
  153. if setup_state.add_url_rule.__name__ != '_blueprint_setup_add_url_rule_patch':
  154. setup_state._original_add_url_rule = setup_state.add_url_rule
  155. setup_state.add_url_rule = MethodType(Api._blueprint_setup_add_url_rule_patch,
  156. setup_state)
  157. if not setup_state.first_registration:
  158. raise ValueError('flask-restful blueprints can only be registered once.')
  159. self._init_app(setup_state.app)
  160. def _init_app(self, app):
  161. """Perform initialization actions with the given :class:`flask.Flask`
  162. object.
  163. :param app: The flask application object
  164. :type app: flask.Flask
  165. """
  166. app.handle_exception = partial(self.error_router, app.handle_exception)
  167. app.handle_user_exception = partial(self.error_router, app.handle_user_exception)
  168. if len(self.resources) > 0:
  169. for resource, urls, kwargs in self.resources:
  170. self._register_view(app, resource, *urls, **kwargs)
  171. def owns_endpoint(self, endpoint):
  172. """Tests if an endpoint name (not path) belongs to this Api. Takes
  173. in to account the Blueprint name part of the endpoint name.
  174. :param endpoint: The name of the endpoint being checked
  175. :return: bool
  176. """
  177. if self.blueprint:
  178. if endpoint.startswith(self.blueprint.name):
  179. endpoint = endpoint.split(self.blueprint.name + '.', 1)[-1]
  180. else:
  181. return False
  182. return endpoint in self.endpoints
  183. def _should_use_fr_error_handler(self):
  184. """ Determine if error should be handled with FR or default Flask
  185. The goal is to return Flask error handlers for non-FR-related routes,
  186. and FR errors (with the correct media type) for FR endpoints. This
  187. method currently handles 404 and 405 errors.
  188. :return: bool
  189. """
  190. adapter = current_app.create_url_adapter(request)
  191. try:
  192. adapter.match()
  193. except MethodNotAllowed as e:
  194. # Check if the other HTTP methods at this url would hit the Api
  195. valid_route_method = e.valid_methods[0]
  196. rule, _ = adapter.match(method=valid_route_method, return_rule=True)
  197. return self.owns_endpoint(rule.endpoint)
  198. except NotFound:
  199. return self.catch_all_404s
  200. except:
  201. # Werkzeug throws other kinds of exceptions, such as Redirect
  202. pass
  203. def _has_fr_route(self):
  204. """Encapsulating the rules for whether the request was to a Flask endpoint"""
  205. # 404's, 405's, which might not have a url_rule
  206. if self._should_use_fr_error_handler():
  207. return True
  208. # for all other errors, just check if FR dispatched the route
  209. if not request.url_rule:
  210. return False
  211. return self.owns_endpoint(request.url_rule.endpoint)
  212. def error_router(self, original_handler, e):
  213. """This function decides whether the error occured in a flask-restful
  214. endpoint or not. If it happened in a flask-restful endpoint, our
  215. handler will be dispatched. If it happened in an unrelated view, the
  216. app's original error handler will be dispatched.
  217. In the event that the error occurred in a flask-restful endpoint but
  218. the local handler can't resolve the situation, the router will fall
  219. back onto the original_handler as last resort.
  220. :param original_handler: the original Flask error handler for the app
  221. :type original_handler: function
  222. :param e: the exception raised while handling the request
  223. :type e: Exception
  224. """
  225. if self._has_fr_route():
  226. try:
  227. return self.handle_error(e)
  228. except Exception:
  229. pass # Fall through to original handler
  230. return original_handler(e)
  231. def handle_error(self, e):
  232. """Error handler for the API transforms a raised exception into a Flask
  233. response, with the appropriate HTTP status code and body.
  234. :param e: the raised Exception object
  235. :type e: Exception
  236. """
  237. got_request_exception.send(current_app._get_current_object(), exception=e)
  238. if not isinstance(e, HTTPException) and current_app.propagate_exceptions:
  239. exc_type, exc_value, tb = sys.exc_info()
  240. if exc_value is e:
  241. raise
  242. else:
  243. raise e
  244. headers = Headers()
  245. if isinstance(e, HTTPException):
  246. if e.response is not None:
  247. # If HTTPException is initialized with a response, then return e.get_response().
  248. # This prevents specified error response from being overridden.
  249. # eg. HTTPException(response=Response("Hello World"))
  250. resp = e.get_response()
  251. return resp
  252. code = e.code
  253. default_data = {
  254. 'message': getattr(e, 'description', http_status_message(code))
  255. }
  256. headers = e.get_response().headers
  257. else:
  258. code = 500
  259. default_data = {
  260. 'message': http_status_message(code),
  261. }
  262. # Werkzeug exceptions generate a content-length header which is added
  263. # to the response in addition to the actual content-length header
  264. # https://github.com/flask-restful/flask-restful/issues/534
  265. remove_headers = ('Content-Length',)
  266. for header in remove_headers:
  267. headers.pop(header, None)
  268. data = getattr(e, 'data', default_data)
  269. if code and code >= 500:
  270. exc_info = sys.exc_info()
  271. if exc_info[1] is None:
  272. exc_info = None
  273. current_app.log_exception(exc_info)
  274. error_cls_name = type(e).__name__
  275. if error_cls_name in self.errors:
  276. custom_data = self.errors.get(error_cls_name, {})
  277. code = custom_data.get('status', 500)
  278. data.update(custom_data)
  279. if code == 406 and self.default_mediatype is None:
  280. # if we are handling NotAcceptable (406), make sure that
  281. # make_response uses a representation we support as the
  282. # default mediatype (so that make_response doesn't throw
  283. # another NotAcceptable error).
  284. supported_mediatypes = list(self.representations.keys())
  285. fallback_mediatype = supported_mediatypes[0] if supported_mediatypes else "text/plain"
  286. resp = self.make_response(
  287. data,
  288. code,
  289. headers,
  290. fallback_mediatype = fallback_mediatype
  291. )
  292. else:
  293. resp = self.make_response(data, code, headers)
  294. if code == 401:
  295. resp = self.unauthorized(resp)
  296. return resp
  297. def mediatypes_method(self):
  298. """Return a method that returns a list of mediatypes
  299. """
  300. return lambda resource_cls: self.mediatypes() + [self.default_mediatype]
  301. def add_resource(self, resource, *urls, **kwargs):
  302. """Adds a resource to the api.
  303. :param resource: the class name of your resource
  304. :type resource: :class:`Type[Resource]`
  305. :param urls: one or more url routes to match for the resource, standard
  306. flask routing rules apply. Any url variables will be
  307. passed to the resource method as args.
  308. :type urls: str
  309. :param endpoint: endpoint name (defaults to :meth:`Resource.__name__.lower`
  310. Can be used to reference this route in :class:`fields.Url` fields
  311. :type endpoint: str
  312. :param resource_class_args: args to be forwarded to the constructor of
  313. the resource.
  314. :type resource_class_args: tuple
  315. :param resource_class_kwargs: kwargs to be forwarded to the constructor
  316. of the resource.
  317. :type resource_class_kwargs: dict
  318. Additional keyword arguments not specified above will be passed as-is
  319. to :meth:`flask.Flask.add_url_rule`.
  320. Examples::
  321. api.add_resource(HelloWorld, '/', '/hello')
  322. api.add_resource(Foo, '/foo', endpoint="foo")
  323. api.add_resource(FooSpecial, '/special/foo', endpoint="foo")
  324. """
  325. if self.app is not None:
  326. self._register_view(self.app, resource, *urls, **kwargs)
  327. else:
  328. self.resources.append((resource, urls, kwargs))
  329. def resource(self, *urls, **kwargs):
  330. """Wraps a :class:`~flask_restful.Resource` class, adding it to the
  331. api. Parameters are the same as :meth:`~flask_restful.Api.add_resource`.
  332. Example::
  333. app = Flask(__name__)
  334. api = restful.Api(app)
  335. @api.resource('/foo')
  336. class Foo(Resource):
  337. def get(self):
  338. return 'Hello, World!'
  339. """
  340. def decorator(cls):
  341. self.add_resource(cls, *urls, **kwargs)
  342. return cls
  343. return decorator
  344. def _register_view(self, app, resource, *urls, **kwargs):
  345. endpoint = kwargs.pop('endpoint', None) or resource.__name__.lower()
  346. self.endpoints.add(endpoint)
  347. resource_class_args = kwargs.pop('resource_class_args', ())
  348. resource_class_kwargs = kwargs.pop('resource_class_kwargs', {})
  349. # NOTE: 'view_functions' is cleaned up from Blueprint class in Flask 1.0
  350. if endpoint in getattr(app, 'view_functions', {}):
  351. previous_view_class = app.view_functions[endpoint].__dict__['view_class']
  352. # if you override the endpoint with a different class, avoid the collision by raising an exception
  353. if previous_view_class != resource:
  354. raise ValueError('This endpoint (%s) is already set to the class %s.' % (endpoint, previous_view_class.__name__))
  355. resource.mediatypes = self.mediatypes_method() # Hacky
  356. resource.endpoint = endpoint
  357. resource_func = self.output(resource.as_view(endpoint, *resource_class_args,
  358. **resource_class_kwargs))
  359. for decorator in self.decorators:
  360. resource_func = decorator(resource_func)
  361. for url in urls:
  362. # If this Api has a blueprint
  363. if self.blueprint:
  364. # And this Api has been setup
  365. if self.blueprint_setup:
  366. # Set the rule to a string directly, as the blueprint is already
  367. # set up.
  368. self.blueprint_setup.add_url_rule(url, view_func=resource_func, **kwargs)
  369. continue
  370. else:
  371. # Set the rule to a function that expects the blueprint prefix
  372. # to construct the final url. Allows deferment of url finalization
  373. # in the case that the associated Blueprint has not yet been
  374. # registered to an application, so we can wait for the registration
  375. # prefix
  376. rule = partial(self._complete_url, url)
  377. else:
  378. # If we've got no Blueprint, just build a url with no prefix
  379. rule = self._complete_url(url, '')
  380. # Add the url to the application or blueprint
  381. app.add_url_rule(rule, view_func=resource_func, **kwargs)
  382. def output(self, resource):
  383. """Wraps a resource (as a flask view function), for cases where the
  384. resource does not directly return a response object
  385. :param resource: The resource as a flask view function
  386. """
  387. @wraps(resource)
  388. def wrapper(*args, **kwargs):
  389. resp = resource(*args, **kwargs)
  390. if isinstance(resp, ResponseBase): # There may be a better way to test
  391. return resp
  392. data, code, headers = unpack(resp)
  393. return self.make_response(data, code, headers=headers)
  394. return wrapper
  395. def url_for(self, resource, **values):
  396. """Generates a URL to the given resource.
  397. Works like :func:`flask.url_for`."""
  398. endpoint = resource.endpoint
  399. if self.blueprint:
  400. endpoint = '{0}.{1}'.format(self.blueprint.name, endpoint)
  401. return url_for(endpoint, **values)
  402. def make_response(self, data, *args, **kwargs):
  403. """Looks up the representation transformer for the requested media
  404. type, invoking the transformer to create a response object. This
  405. defaults to default_mediatype if no transformer is found for the
  406. requested mediatype. If default_mediatype is None, a 406 Not
  407. Acceptable response will be sent as per RFC 2616 section 14.1
  408. :param data: Python object containing response data to be transformed
  409. """
  410. default_mediatype = kwargs.pop('fallback_mediatype', None) or self.default_mediatype
  411. mediatype = request.accept_mimetypes.best_match(
  412. self.representations,
  413. default=default_mediatype,
  414. )
  415. if mediatype is None:
  416. raise NotAcceptable()
  417. if mediatype in self.representations:
  418. resp = self.representations[mediatype](data, *args, **kwargs)
  419. resp.headers['Content-Type'] = mediatype
  420. return resp
  421. elif mediatype == 'text/plain':
  422. resp = original_flask_make_response(str(data), *args, **kwargs)
  423. resp.headers['Content-Type'] = 'text/plain'
  424. return resp
  425. else:
  426. raise InternalServerError()
  427. def mediatypes(self):
  428. """Returns a list of requested mediatypes sent in the Accept header"""
  429. return [h for h, q in sorted(request.accept_mimetypes,
  430. key=operator.itemgetter(1), reverse=True)]
  431. def representation(self, mediatype):
  432. """Allows additional representation transformers to be declared for the
  433. api. Transformers are functions that must be decorated with this
  434. method, passing the mediatype the transformer represents. Three
  435. arguments are passed to the transformer:
  436. * The data to be represented in the response body
  437. * The http status code
  438. * A dictionary of headers
  439. The transformer should convert the data appropriately for the mediatype
  440. and return a Flask response object.
  441. Ex::
  442. @api.representation('application/xml')
  443. def xml(data, code, headers):
  444. resp = make_response(convert_data_to_xml(data), code)
  445. resp.headers.extend(headers)
  446. return resp
  447. """
  448. def wrapper(func):
  449. self.representations[mediatype] = func
  450. return func
  451. return wrapper
  452. def unauthorized(self, response):
  453. """ Given a response, change it to ask for credentials """
  454. if self.serve_challenge_on_401:
  455. realm = current_app.config.get("HTTP_BASIC_AUTH_REALM", "flask-restful")
  456. challenge = u"{0} realm=\"{1}\"".format("Basic", realm)
  457. response.headers['WWW-Authenticate'] = challenge
  458. return response
  459. class Resource(MethodView):
  460. """
  461. Represents an abstract RESTful resource. Concrete resources should
  462. extend from this class and expose methods for each supported HTTP
  463. method. If a resource is invoked with an unsupported HTTP method,
  464. the API will return a response with status 405 Method Not Allowed.
  465. Otherwise the appropriate method is called and passed all arguments
  466. from the url rule used when adding the resource to an Api instance. See
  467. :meth:`~flask_restful.Api.add_resource` for details.
  468. """
  469. representations = None
  470. method_decorators = []
  471. def dispatch_request(self, *args, **kwargs):
  472. # Taken from flask
  473. #noinspection PyUnresolvedReferences
  474. meth = getattr(self, request.method.lower(), None)
  475. if meth is None and request.method == 'HEAD':
  476. meth = getattr(self, 'get', None)
  477. assert meth is not None, 'Unimplemented method %r' % request.method
  478. if isinstance(self.method_decorators, Mapping):
  479. decorators = self.method_decorators.get(request.method.lower(), [])
  480. else:
  481. decorators = self.method_decorators
  482. for decorator in decorators:
  483. meth = decorator(meth)
  484. resp = meth(*args, **kwargs)
  485. if isinstance(resp, ResponseBase): # There may be a better way to test
  486. return resp
  487. representations = self.representations or OrderedDict()
  488. #noinspection PyUnresolvedReferences
  489. mediatype = request.accept_mimetypes.best_match(representations, default=None)
  490. if mediatype in representations:
  491. data, code, headers = unpack(resp)
  492. resp = representations[mediatype](data, code, headers)
  493. resp.headers['Content-Type'] = mediatype
  494. return resp
  495. return resp
  496. def marshal(data, fields, envelope=None):
  497. """Takes raw data (in the form of a dict, list, object) and a dict of
  498. fields to output and filters the data based on those fields.
  499. :param data: the actual object(s) from which the fields are taken from
  500. :param fields: a dict of whose keys will make up the final serialized
  501. response output
  502. :param envelope: optional key that will be used to envelop the serialized
  503. response
  504. >>> from flask_restful import fields, marshal
  505. >>> data = { 'a': 100, 'b': 'foo' }
  506. >>> mfields = { 'a': fields.Raw }
  507. >>> marshal(data, mfields)
  508. OrderedDict([('a', 100)])
  509. >>> marshal(data, mfields, envelope='data')
  510. OrderedDict([('data', OrderedDict([('a', 100)]))])
  511. """
  512. def make(cls):
  513. if isinstance(cls, type):
  514. return cls()
  515. return cls
  516. if isinstance(data, (list, tuple)):
  517. return (OrderedDict([(envelope, [marshal(d, fields) for d in data])])
  518. if envelope else [marshal(d, fields) for d in data])
  519. items = ((k, marshal(data, v) if isinstance(v, dict)
  520. else make(v).output(k, data))
  521. for k, v in fields.items())
  522. return OrderedDict([(envelope, OrderedDict(items))]) if envelope else OrderedDict(items)
  523. class marshal_with(object):
  524. """A decorator that apply marshalling to the return values of your methods.
  525. >>> from flask_restful import fields, marshal_with
  526. >>> mfields = { 'a': fields.Raw }
  527. >>> @marshal_with(mfields)
  528. ... def get():
  529. ... return { 'a': 100, 'b': 'foo' }
  530. ...
  531. ...
  532. >>> get()
  533. OrderedDict([('a', 100)])
  534. >>> @marshal_with(mfields, envelope='data')
  535. ... def get():
  536. ... return { 'a': 100, 'b': 'foo' }
  537. ...
  538. ...
  539. >>> get()
  540. OrderedDict([('data', OrderedDict([('a', 100)]))])
  541. see :meth:`flask_restful.marshal`
  542. """
  543. def __init__(self, fields, envelope=None):
  544. """
  545. :param fields: a dict of whose keys will make up the final
  546. serialized response output
  547. :param envelope: optional key that will be used to envelop the serialized
  548. response
  549. """
  550. self.fields = fields
  551. self.envelope = envelope
  552. def __call__(self, f):
  553. @wraps(f)
  554. def wrapper(*args, **kwargs):
  555. resp = f(*args, **kwargs)
  556. if isinstance(resp, tuple):
  557. data, code, headers = unpack(resp)
  558. return marshal(data, self.fields, self.envelope), code, headers
  559. else:
  560. return marshal(resp, self.fields, self.envelope)
  561. return wrapper
  562. class marshal_with_field(object):
  563. """
  564. A decorator that formats the return values of your methods with a single field.
  565. >>> from flask_restful import marshal_with_field, fields
  566. >>> @marshal_with_field(fields.List(fields.Integer))
  567. ... def get():
  568. ... return ['1', 2, 3.0]
  569. ...
  570. >>> get()
  571. [1, 2, 3]
  572. see :meth:`flask_restful.marshal_with`
  573. """
  574. def __init__(self, field):
  575. """
  576. :param field: a single field with which to marshal the output.
  577. """
  578. if isinstance(field, type):
  579. self.field = field()
  580. else:
  581. self.field = field
  582. def __call__(self, f):
  583. @wraps(f)
  584. def wrapper(*args, **kwargs):
  585. resp = f(*args, **kwargs)
  586. if isinstance(resp, tuple):
  587. data, code, headers = unpack(resp)
  588. return self.field.format(data), code, headers
  589. return self.field.format(resp)
  590. return wrapper