
import json

from flask import Response, jsonify


class ContextResponse(Response):
    """
    This class adds the "@context" parameter to JSON responses before
    they're sent to the user.

    For an explanation of how this works/is used, read
    https://blog.miguelgrinberg.com/post/customizing-the-flask-response-class
    """
    context_url = "https://api.wellcomecollection.org/storage/v1/context.json"

    def __init__(self, response, **kwargs):
        # Here we unmarshal the response as provided by Flask-RESTPlus, add
        # the @context parameter, then repack it.
        rv = json.loads(response)

        # The @context may already be provided if we've been through the
        # force_type method below.
        if "@context" in rv:
            return super(ContextResponse, self).__init__(response, **kwargs)
        else:
            rv["@context"] = self.context_url
            return super(ContextResponse, self).__init__(json.dumps(rv), **kwargs)

    @classmethod
    def force_type(cls, rv, environ=None):
        # All of our endpoints should be returning a dictionary to be
        # serialised as JSON.
        assert isinstance(rv, dict)

        assert "@context" not in rv, rv
        rv["@context"] = cls.context_url

        return super(ContextResponse, cls).force_type(jsonify(rv), environ)


import json

from flask import Response, jsonify
from werkzeug.wsgi import ClosingIterator


class ContextResponse(Response):
    """
    This class adds the "@context" parameter to JSON responses before
    they're sent to the user.

    For an explanation of how this works/is used, read
    https://blog.miguelgrinberg.com/post/customizing-the-flask-response-class
    """
    context_url = "https://api.wellcomecollection.org/storage/v1/context.json"

    def __init__(self, response, *args, **kwargs):
        """
        Unmarshal the response as provided by Flask-RESTPlus, add the
        @context parameter, then repack it.
        """
        if isinstance(response, ClosingIterator):
            response = b''.join([char for char in response])

        rv = json.loads(response)

        # The @context may already be provided if we've been through the
        # force_type method below.
        if "@context" in rv:
            return super(ContextResponse, self).__init__(response, **kwargs)
        else:
            rv["@context"] = self.context_url
            json_string = json.dumps(rv)
            return super(ContextResponse, self).__init__(json_string, **kwargs)

    @classmethod
    def force_type(cls, rv, environ=None):
        # All of our endpoints should be returning a dictionary to be
        # serialised as JSON.
        assert isinstance(rv, dict)

        assert "@context" not in rv, rv
        rv["@context"] = cls.context_url

        return super(ContextResponse, cls).force_type(jsonify(rv), environ)

Handle a Werkzeug ClosingIterator (as exposed by the tests)
Handle a Werkzeug ClosingIterator (as exposed by the tests)

  
  import json
  
  from flask import Response, jsonify
+ from werkzeug.wsgi import ClosingIterator
  
  
  class ContextResponse(Response):
      """
      This class adds the "@context" parameter to JSON responses before
      they're sent to the user.
  
      For an explanation of how this works/is used, read
      https://blog.miguelgrinberg.com/post/customizing-the-flask-response-class
      """
      context_url = "https://api.wellcomecollection.org/storage/v1/context.json"
  
-     def __init__(self, response, **kwargs):
+     def __init__(self, response, *args, **kwargs):
+         """
-         # Here we unmarshal the response as provided by Flask-RESTPlus, add
+         Unmarshal the response as provided by Flask-RESTPlus, add the
-         # the @context parameter, then repack it.
+         @context parameter, then repack it.
+         """
+         if isinstance(response, ClosingIterator):
+             response = b''.join([char for char in response])
+ 
          rv = json.loads(response)
  
          # The @context may already be provided if we've been through the
          # force_type method below.
          if "@context" in rv:
              return super(ContextResponse, self).__init__(response, **kwargs)
          else:
              rv["@context"] = self.context_url
+             json_string = json.dumps(rv)
-             return super(ContextResponse, self).__init__(json.dumps(rv), **kwargs)
+             return super(ContextResponse, self).__init__(json_string, **kwargs)
  
      @classmethod
      def force_type(cls, rv, environ=None):
          # All of our endpoints should be returning a dictionary to be
          # serialised as JSON.
          assert isinstance(rv, dict)
  
          assert "@context" not in rv, rv
          rv["@context"] = cls.context_url
  
          return super(ContextResponse, cls).force_type(jsonify(rv), environ)
  
Handle a Werkzeug ClosingIterator (as exposed by the tests)
## Code Before:

import json

from flask import Response, jsonify


class ContextResponse(Response):
    """
    This class adds the "@context" parameter to JSON responses before
    they're sent to the user.

    For an explanation of how this works/is used, read
    https://blog.miguelgrinberg.com/post/customizing-the-flask-response-class
    """
    context_url = "https://api.wellcomecollection.org/storage/v1/context.json"

    def __init__(self, response, **kwargs):
        # Here we unmarshal the response as provided by Flask-RESTPlus, add
        # the @context parameter, then repack it.
        rv = json.loads(response)

        # The @context may already be provided if we've been through the
        # force_type method below.
        if "@context" in rv:
            return super(ContextResponse, self).__init__(response, **kwargs)
        else:
            rv["@context"] = self.context_url
            return super(ContextResponse, self).__init__(json.dumps(rv), **kwargs)

    @classmethod
    def force_type(cls, rv, environ=None):
        # All of our endpoints should be returning a dictionary to be
        # serialised as JSON.
        assert isinstance(rv, dict)

        assert "@context" not in rv, rv
        rv["@context"] = cls.context_url

        return super(ContextResponse, cls).force_type(jsonify(rv), environ)

## Instruction:
Handle a Werkzeug ClosingIterator (as exposed by the tests)
## Code After:

import json

from flask import Response, jsonify
from werkzeug.wsgi import ClosingIterator


class ContextResponse(Response):
    """
    This class adds the "@context" parameter to JSON responses before
    they're sent to the user.

    For an explanation of how this works/is used, read
    https://blog.miguelgrinberg.com/post/customizing-the-flask-response-class
    """
    context_url = "https://api.wellcomecollection.org/storage/v1/context.json"

    def __init__(self, response, *args, **kwargs):
        """
        Unmarshal the response as provided by Flask-RESTPlus, add the
        @context parameter, then repack it.
        """
        if isinstance(response, ClosingIterator):
            response = b''.join([char for char in response])

        rv = json.loads(response)

        # The @context may already be provided if we've been through the
        # force_type method below.
        if "@context" in rv:
            return super(ContextResponse, self).__init__(response, **kwargs)
        else:
            rv["@context"] = self.context_url
            json_string = json.dumps(rv)
            return super(ContextResponse, self).__init__(json_string, **kwargs)

    @classmethod
    def force_type(cls, rv, environ=None):
        # All of our endpoints should be returning a dictionary to be
        # serialised as JSON.
        assert isinstance(rv, dict)

        assert "@context" not in rv, rv
        rv["@context"] = cls.context_url

        return super(ContextResponse, cls).force_type(jsonify(rv), environ)
