def test_comp(typestr, comptype, index, iter_vars_names,
              not_called_msg, insufficient_ifs_msg, incorrect_iter_vars_msg,
              comp_iter, ifs, key=None, body=None, value=None,
              rep=None, state=None):

    MSG_INCORRECT_ITER_VARS = "Have you used the correct iterator variables?"
    MSG_INCORRECT_NUM_ITER_VARS = "Have you used {{num_vars}} iterator variables?"
    MSG_INSUFFICIENT_IFS = "Have you used {{sol_len}} ifs?"

    # make sure other messages are set to default if None
    if insufficient_ifs_msg is None:
        insufficient_ifs_msg = MSG_INSUFFICIENT_IFS

    # get comprehension
    child = check_node(comptype, index-1, typestr, missing_msg=not_called_msg, state=state)

    # test comprehension iter and its variable names (or number of variables)
    if comp_iter: multi(comp_iter, state=check_part("iter", "iterable part", state=child))

    # test iterator variables
    default_msg = MSG_INCORRECT_ITER_VARS if iter_vars_names else MSG_INCORRECT_NUM_ITER_VARS
    has_context(incorrect_iter_vars_msg or default_msg, iter_vars_names, state=child)

    # test the main expressions.
    if body:   multi(body,  state=check_part("body", "body", state=child))        # list and gen comp
    if key:    multi(key,   state=check_part("key", "key part",  state=child))    # dict comp
    if value:  multi(value, state=check_part("value", "value part", state=child)) # ""

    # test a list of ifs. each entry corresponds to a filter in the comprehension.
    for i, if_test in enumerate(ifs or []):
        # test that ifs are same length
        has_equal_part_len('ifs', insufficient_ifs_msg, state=child)
        # test individual ifs
        multi(if_test, state=check_part_index("ifs", i, utils.get_ord(i+1) + " if", state=child))
def test_with(
        index,
        context_vals=False,  # whether to check number of context vals
        context_tests=None,  # check on context expressions
        body=None,
        undefined_msg=None,
        context_vals_len_msg=None,
        context_vals_msg=None,
        expand_message=True,
        state=None):
    """Test a with statement.
with open_file('...') as bla:

    [ open_file('...').__enter__() ]


with open_file('...') as file:
    [ ]

    """

    check_with = partial(check_node,
                         'withs',
                         index - 1,
                         "`with` statement",
                         MSG_MISSING,
                         state=state)

    child = check_with(MSG_PREPEND if expand_message else "")
    child2 = check_with(MSG_PREPEND2 if expand_message else "")

    if context_vals:
        # test context var names ----
        has_context(incorrect_msg=context_vals_msg or MSG_CTXT_NAMES,
                    exact_names=True,
                    state=child)

        # test num context vars ----
        has_equal_part_len('context', MSG_NUM_CTXT, state=child)

    # Context sub tests ----
    if context_tests and not isinstance(context_tests, list):
        context_tests = [context_tests]

    for i, context_test in enumerate(context_tests or []):
        # partial the substate check, because the function uses two prepended messages
        check_context = partial(check_part_index, 'context', i,
                                "%s context" % utils.get_ord(i + 1),
                                MSG_NUM_CTXT2)

        check_context(state=child)  # test exist

        ctxt_state = check_context(state=child2)  # sub tests
        multi(context_test, state=ctxt_state)

    # Body sub tests ----
    if body is not None:
        body_state = check_part('body', 'body', state=child2)

        with_context(body, state=body_state)
def test_comp(typestr,
              comptype,
              index,
              iter_vars_names,
              not_called_msg,
              insufficient_ifs_msg,
              incorrect_iter_vars_msg,
              comp_iter,
              ifs,
              key=None,
              body=None,
              value=None,
              expand_message=True,
              rep=None,
              state=None):

    # if true, set expand_message to default (for backwards compatibility)
    expand_message = MSG_PREPEND if expand_message is True else (expand_message
                                                                 or "")
    # make sure other messages are set to default if None
    if insufficient_ifs_msg is None:
        insufficient_ifs_msg = MSG_INSUFFICIENT_IFS
    if not_called_msg is None: not_called_msg = MSG_NOT_CALLED

    # TODO MSG: function was not consistent with prepending, so use state w/o expand_message
    quiet_state = check_node(comptype, index - 1, typestr, not_called_msg, "",
                             state)

    # get comprehension
    state = check_node(comptype, index - 1, typestr, not_called_msg,
                       expand_message, state)

    # test comprehension iter and its variable names (or number of variables)
    if comp_iter:
        multi(comp_iter, state=check_part("iter", "iterable part", state))

    # test iterator variables
    default_msg = MSG_INCORRECT_ITER_VARS if iter_vars_names else MSG_INCORRECT_NUM_ITER_VARS
    has_context(incorrect_iter_vars_msg or default_msg,
                iter_vars_names,
                state=quiet_state)

    # test the main expressions.
    if body:
        multi(body, state=check_part("body", "body",
                                     state))  # list and gen comp
    if key: multi(key, state=check_part("key", "key part", state))  # dict comp
    if value:
        multi(value, state=check_part("value", "value part", state))  # ""

    # test a list of ifs. each entry corresponds to a filter in the comprehension.
    for i, if_test in enumerate(ifs or []):
        # test that ifs are same length
        has_equal_part_len('ifs', insufficient_ifs_msg, state=quiet_state)
        # test individual ifs
        multi(if_test,
              state=check_part_index("ifs",
                                     i,
                                     get_ord(i + 1) + " if",
                                     state=state))
def test_with(index,
              context_vals=False, # whether to check number of context vals
              context_tests=None, # check on context expressions
              body=None,
              undefined_msg=None,
              context_vals_len_msg=None,
              context_vals_msg=None,
              state=None):
    """Test a with statement.
with open_file('...') as bla:

    [ open_file('...').__enter__() ]


with open_file('...') as file:
    [ ]

    """

    MSG_NUM_CTXT = "Make sure to use the correct number of context variables. It seems you defined too many."
    MSG_NUM_CTXT2 = "Make sure to use the correct number of context variables. It seems you defined too little."
    MSG_CTXT_NAMES = "Make sure to use the correct context variable names. Was expecting `{{sol_vars}}` but got `{{stu_vars}}`."

    check_with = partial(check_node, 'withs', index-1, "{{ordinal}} `with` statement", state=state)

    child =  check_with()
    child2 = check_with()

    if context_vals:
        # test context var names ----
        has_context(incorrect_msg=context_vals_msg or MSG_CTXT_NAMES, exact_names = True, state=child)

        # test num context vars ----
        has_equal_part_len('context', MSG_NUM_CTXT, state=child)
        
    
    # Context sub tests ----
    if context_tests and not isinstance(context_tests, list): context_tests = [context_tests]

    for i, context_test in enumerate(context_tests or []):
        # partial the substate check, because the function uses two prepended messages
        check_context = partial(check_part_index, 'context', i, "%s context"%utils.get_ord(i+1), missing_msg=MSG_NUM_CTXT2)

        check_context(state=child)                   # test exist

        ctxt_state = check_context(state=child2)     # sub tests
        multi(context_test, state=ctxt_state)
    
    # Body sub tests ----
    if body is not None:
        body_state = check_part('body', 'body', state=child2)

        with_context(body, state=body_state)