def set_context(state, *args, **kwargs):
    """Update context values for student and solution environments.
    
    When ``has_equal_x()`` is used after this, the context values (in ``for`` loops and function definitions, for example)
    will have the values specified through his function. It is the function equivalent of the ``context_vals`` argument of
    the ``has_equal_x()`` functions.

    - Note 1: excess args and unmatched kwargs will be unused in the student environment.
    - Note 2: When you try to set context values that don't match any target variables in the solution code,
      ``set_context()`` raises an exception that lists the ones available.
    - Note 3: positional arguments are more robust to the student using different names for context values.
    - Note 4: You have to specify arguments either by position, either by name. A combination is not possible.

    :Example:

        Solution code::

            total = 0
            for i in range(10):
                print(i ** 2)

        Student submission that will pass (different iterator, different calculation)::

            total = 0
            for j in range(10):
                print(j * j)

        SCT::

            # set_context is robust against different names of context values.
            Ex().check_for_loop().check_body().multi(
                set_context(1).has_equal_output(),
                set_context(2).has_equal_output(),
                set_context(3).has_equal_output()
            )

            # equivalent SCT, by setting context_vals in has_equal_output()
            Ex().check_for_loop().check_body().\\
                multi([s.has_equal_output(context_vals=[i]) for i in range(1, 4)])

    """

    stu_crnt = state.student_context.context
    sol_crnt = state.solution_context.context

    # for now, you can't specify both
    if len(args) > 0 and len(kwargs) > 0:
        raise InstructorError(
            "In `set_context()`, specify arguments either by position, either by name."
        )

    # set args specified by pos -----------------------------------------------
    if args:
        # stop if too many pos args for solution
        if len(args) > len(sol_crnt):
            raise InstructorError(
                "Too many positional args. There are {} context vals, but tried to set {}".format(
                    len(sol_crnt), len(args)
                )
            )
        # set pos args
        upd_sol = sol_crnt.update(dict(zip(sol_crnt.keys(), args)))
        upd_stu = stu_crnt.update(dict(zip(stu_crnt.keys(), args)))
    else:
        upd_sol = sol_crnt
        upd_stu = stu_crnt

    # set args specified by keyword -------------------------------------------
    if kwargs:
        # stop if keywords don't match with solution
        if set(kwargs) - set(upd_sol):
            raise InstructorError(
                "`set_context()` failed: context val names are {}, but you tried to set {}.".format(
                    upd_sol or "missing", sorted(list(kwargs.keys()))
                )
            )
        out_sol = upd_sol.update(kwargs)
        # need to match keys in kwargs with corresponding keys in stu context
        # in case they used, e.g., different loop variable names
        match_keys = dict(zip(sol_crnt.keys(), stu_crnt.keys()))
        out_stu = upd_stu.update(
            {match_keys[k]: v for k, v in kwargs.items() if k in match_keys}
        )
    else:
        out_sol = upd_sol
        out_stu = upd_stu

    return state.to_child(
        student_context=out_stu, solution_context=out_sol, highlight=state.highlight
    )
Update context values for student and solution environments.
    
    When ``has_equal_x()`` is used after this, the context values (in ``for`` loops and function definitions, for example)
    will have the values specified through his function. It is the function equivalent of the ``context_vals`` argument of
    the ``has_equal_x()`` functions.

    - Note 1: excess args and unmatched kwargs will be unused in the student environment.
    - Note 2: When you try to set context values that don't match any target variables in the solution code,
      ``set_context()`` raises an exception that lists the ones available.
    - Note 3: positional arguments are more robust to the student using different names for context values.
    - Note 4: You have to specify arguments either by position, either by name. A combination is not possible.

    :Example:

        Solution code::

            total = 0
            for i in range(10):
                print(i ** 2)

        Student submission that will pass (different iterator, different calculation)::

            total = 0
            for j in range(10):
                print(j * j)

        SCT::

            # set_context is robust against different names of context values.
            Ex().check_for_loop().check_body().multi(
                set_context(1).has_equal_output(),
                set_context(2).has_equal_output(),
                set_context(3).has_equal_output()
            )

            # equivalent SCT, by setting context_vals in has_equal_output()
            Ex().check_for_loop().check_body().\\
                multi([s.has_equal_output(context_vals=[i]) for i in range(1, 4)])